Kernel: cosine_sim (Level 1, category: loss)
Input shapes: [[2048, 1024], [2048, 1024]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 413.7239 ms; JAX XLA compiled runtime: 0.1271 ms; speedup vs compiled baseline: 0.0003

JAX reference implementation:
@jax.jit
def jax_cosine_sim(x: jax.Array, y: jax.Array) -> jax.Array:
    dot = jnp.sum(x * y, axis=-1)
    norm_x = jnp.sqrt(jnp.sum(x * x, axis=-1))
    norm_y = jnp.sqrt(jnp.sum(y * y, axis=-1))
    return dot / (norm_x * norm_y + 1e-8)

Original Pallas kernel (TPU-oriented): """Level 1: Row-wise cosine similarity via Pallas.

Provenance: standard similarity metric for embeddings and retrieval
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _cosine_sim_kernel(x_ref, y_ref, o_ref):
    x = x_ref[...]
    y = y_ref[...]
    dot = jnp.sum(x * y, axis=-1)
    norm_x = jnp.sqrt(jnp.sum(x * x, axis=-1))
    norm_y = jnp.sqrt(jnp.sum(y * y, axis=-1))
    o_ref[...] = dot / (norm_x * norm_y + 1e-8)


def pallas_cosine_sim(x: jax.Array, y: jax.Array) -> jax.Array:
    n_rows = x.shape[0]
    n_cols = x.shape[1]
    block_rows = min(256, n_rows)
    grid_size = n_rows // block_rows

    return pl.pallas_call(
        _cosine_sim_kernel,
        out_shape=jax.ShapeDtypeStruct((n_rows,), x.dtype),
        grid=(grid_size,),
        in_specs=[
            pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0)),
            pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0)),
        ],
        out_specs=pl.BlockSpec((block_rows,), lambda i: (i,)),
    )(x, y)


pallas_kernel = pallas_cosine_sim
task_name = "cosine_sim"
input_shapes = [(2048, 1024), (2048, 1024)]
category = "loss"
level = 1

Fixed Pallas kernel (GPU-compatible): """Level 1: Row-wise cosine similarity via Pallas.

Provenance: standard similarity metric for embeddings and retrieval
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _cosine_sim_kernel(x_ref, y_ref, o_ref):
    x = x_ref[...]
    y = y_ref[...]
    dot = jnp.sum(x * y, axis=-1)
    norm_x = jnp.sqrt(jnp.sum(x * x, axis=-1))
    norm_y = jnp.sqrt(jnp.sum(y * y, axis=-1))
    o_ref[...] = dot / (norm_x * norm_y + 1e-8)


def pallas_cosine_sim(x: jax.Array, y: jax.Array) -> jax.Array:
    n_rows = x.shape[0]
    n_cols = x.shape[1]
    block_rows = min(256, n_rows)
    while block_rows * n_cols > 16384 and block_rows > 1:
        block_rows //= 2
    grid_size = n_rows // block_rows

    return pl.pallas_call(
        _cosine_sim_kernel,
        out_shape=jax.ShapeDtypeStruct((n_rows,), x.dtype),
        grid=(grid_size,),
        in_specs=[
            pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0)),
            pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0)),
        ],
        out_specs=pl.BlockSpec((block_rows,), lambda i: (i,)),
    )(x, y)


pallas_kernel = pallas_cosine_sim
task_name = "cosine_sim"
input_shapes = [(2048, 1024), (2048, 1024)]
category = "loss"
level = 1

Diff of GPU fix: --- a/cosine_sim.py
+++ b/cosine_sim.py
@@ -26,6 +26,8 @@
     n_rows = x.shape[0]
     n_cols = x.shape[1]
     block_rows = min(256, n_rows)
+    while block_rows * n_cols > 16384 and block_rows > 1:
+        block_rows //= 2
     grid_size = n_rows // block_rows
 
     return pl.pallas_call(

Jaxpr IR: { lambda ; a:f32[2048,1024] b:f32[2048,1024]. let
    c:f32[2048] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(128,), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=16), Blocked(block_size=1024))), BlockMapping(block_shape=(Blocked(block_size=16), Blocked(block_size=1024))), BlockMapping(block_shape=(Blocked(block_size=16),))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; d:Ref<default>{f32[16,1024]} e:Ref<default>{f32[16,1024]}
          f:Ref<default>{f32[16]}. let
          g:f32[16,1024] <- d[...]
          h:f32[16,1024] <- e[...]
          i:f32[16,1024] = mul g h
          j:f32[16] = reduce_sum[axes=(1,) out_sharding=None] i
          k:f32[16,1024] = mul g g
          l:f32[16] = reduce_sum[axes=(1,) out_sharding=None] k
          m:f32[16] = sqrt l
          n:f32[16,1024] = mul h h
          o:f32[16] = reduce_sum[axes=(1,) out_sharding=None] n
          p:f32[16] = sqrt o
          q:f32[16] = mul m p
          r:f32[16] = add q 9.99999993922529e-09:f32[]
          s:f32[16] = div j r
          t:f32[16], f[...] <- f[...], s
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[2048]),)
    ] a b
  in (c,) }