Kernel: outer_product (Level 1, category: matmul)
Input shapes: [[1024], [1024]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 5730.2626 ms; JAX XLA compiled runtime: 0.0945 ms; speedup vs compiled baseline: 0

JAX reference implementation:
@jax.jit
def jax_outer_product(x: jax.Array, y: jax.Array) -> jax.Array:
    return jnp.outer(x, y)

Original Pallas kernel (TPU-oriented): """Level 1: Outer product via Pallas.

Provenance: jnp.outer, rank-1 update pattern used in Evoformer/AlphaFold
"""

from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L1/outer_product", __doc__)

import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl


def _outer_kernel(x_ref, y_ref, o_ref):
    x = x_ref[...]
    y = y_ref[...]
    o_ref[...] = x[:, None] * y[None, :]


def pallas_outer_product(x: jax.Array, y: jax.Array) -> jax.Array:
    m = x.shape[0]
    n = y.shape[0]

    return pl.pallas_call(
        _outer_kernel,
        out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
        grid=(1,),
        in_specs=[
            pl.BlockSpec((m,), lambda i: (0,)),
            pl.BlockSpec((n,), lambda i: (0,)),
        ],
        out_specs=pl.BlockSpec((m, n), lambda i: (0, 0)),
    )(x, y)


pallas_kernel = pallas_outer_product
task_name = "outer_product"
input_shapes = [(1024,), (1024,)]
category = "matmul"
level = 1

Fixed Pallas kernel (GPU-compatible): """Level 1: Outer product via Pallas.

Provenance: jnp.outer, rank-1 update pattern used in Evoformer/AlphaFold
"""

from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L1/outer_product", __doc__)

import jax
import jax.numpy as jnp
from jax.experimental import pallas as pl


def _outer_kernel(x_ref, y_ref, o_ref):
    x = x_ref[...]
    y = y_ref[...]
    o_ref[...] = x[:, None] * y[None, :]


def pallas_outer_product(x: jax.Array, y: jax.Array) -> jax.Array:
    m = x.shape[0]
    n = y.shape[0]

    return pl.pallas_call(
        _outer_kernel,
        out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
        grid=(1,),
        in_specs=[
            pl.BlockSpec((m,), lambda i: (0,)),
            pl.BlockSpec((n,), lambda i: (0,)),
        ],
        out_specs=pl.BlockSpec((m, n), lambda i: (0, 0)),
    )(x, y)


pallas_kernel = pallas_outer_product
task_name = "outer_product"
input_shapes = [(1024,), (1024,)]
category = "matmul"
level = 1

Jaxpr IR: { lambda ; a:f32[1024] b:f32[1024]. let
    c:f32[1024,1024] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(1,), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=1024),)), BlockMapping(block_shape=(Blocked(block_size=1024),)), BlockMapping(block_shape=(Blocked(block_size=1024), Blocked(block_size=1024)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; d:Ref<default>{f32[1024]} e:Ref<default>{f32[1024]} f:Ref<default>{f32[1024,1024]}. let
          g:f32[1024] <- d[...]
          h:f32[1024] <- e[...]
          i:f32[1024,1] = broadcast_in_dim[broadcast_dimensions=(0,)] g
          j:f32[1,1024] = broadcast_in_dim[broadcast_dimensions=(1,)] h
          k:f32[1024,1024] = mul i j
          l:f32[1024,1024], f[...] <- f[...], k
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[1024,1024]),)
    ] a b
  in (c,) }