EvanOLeary/pallasbench-robust-gpu-a100 huggingface dataset:
Kernel: tanh (Level 1, category: activation)
Input shapes: [[4096, 4096]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 252.8305 ms; JAX XLA compiled runtime: 0.2102 ms; speedup vs compiled baseline: 0.0008

JAX reference implementation:
@jax.jit
def jax_tanh(x: jax.Array) -> jax.Array:
    return jnp.tanh(x)

Jaxpr IR: { lambda ; a:f32[4096,4096]. let
    b:f32[4096,4096] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(32, 32), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128))), BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; c:Ref<default>{f32[128,128]} d:Ref<default>{f32[128,128]}. let
          e:f32[128,128] <- c[...]
          f:f32[128,128] = tanh e
          g:f32[128,128], d[...] <- d[...], f
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[4096,4096]),)
    ] a
  in (b,) }
Kernel: exp (Level 1, category: elementwise)
Input shapes: [[4096, 4096]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 179.5751 ms; JAX XLA compiled runtime: 0.2054 ms; speedup vs compiled baseline: 0.0011

JAX reference implementation:
@jax.jit
def jax_exp(x: jax.Array) -> jax.Array:
    return jnp.exp(x)

Jaxpr IR: { lambda ; a:f32[4096,4096]. let
    b:f32[4096,4096] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(32, 32), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128))), BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; c:Ref<default>{f32[128,128]} d:Ref<default>{f32[128,128]}. let
          e:f32[128,128] <- c[...]
          f:f32[128,128] = exp e
          g:f32[128,128], d[...] <- d[...], f
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[4096,4096]),)
    ] a
  in (b,) }
Kernel: log (Level 1, category: elementwise)
Input shapes: [[4096, 4096]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 430.9604 ms; JAX XLA compiled runtime: 0.2047 ms; speedup vs compiled baseline: 0.0005

JAX reference implementation:
@jax.jit
def jax_log(x: jax.Array) -> jax.Array:
    return jnp.log(x + 1e-7)

Jaxpr IR: { lambda ; a:f32[4096,4096]. let
    b:f32[4096,4096] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(32, 32), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128))), BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; c:Ref<default>{f32[128,128]} d:Ref<default>{f32[128,128]}. let
          e:f32[128,128] <- c[...]
          f:f32[128,128] = add e 1.0000000116860974e-07:f32[]
          g:f32[128,128] = log f
          h:f32[128,128], d[...] <- d[...], g
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[4096,4096]),)
    ] a
  in (b,) }
Kernel: add (Level 1, category: elementwise)
Input shapes: [[4096, 4096], [4096, 4096]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 128.6307 ms; JAX XLA compiled runtime: 0.2465 ms; speedup vs compiled baseline: 0.0019

JAX reference implementation:
@jax.jit
def jax_add(x: jax.Array, y: jax.Array) -> jax.Array:
    return x + y

Jaxpr IR: { lambda ; a:f32[4096,4096] b:f32[4096,4096]. let
    c:f32[4096,4096] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(32, 32), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128))), BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128))), BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; d:Ref<default>{f32[128,128]} e:Ref<default>{f32[128,128]}
          f:Ref<default>{f32[128,128]}. let
          g:f32[128,128] <- d[...]
          h:f32[128,128] <- e[...]
          i:f32[128,128] = add g h
          j:f32[128,128], f[...] <- f[...], i
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[4096,4096]),)
    ] a b
  in (c,) }
Kernel: multiply (Level 1, category: elementwise)
Input shapes: [[4096, 4096], [4096, 4096]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 129.6258 ms; JAX XLA compiled runtime: 0.2472 ms; speedup vs compiled baseline: 0.0019

JAX reference implementation:
@jax.jit
def jax_multiply(x: jax.Array, y: jax.Array) -> jax.Array:
    return x * y

Jaxpr IR: { lambda ; a:f32[4096,4096] b:f32[4096,4096]. let
    c:f32[4096,4096] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(32, 32), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128))), BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128))), BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; d:Ref<default>{f32[128,128]} e:Ref<default>{f32[128,128]}
          f:Ref<default>{f32[128,128]}. let
          g:f32[128,128] <- d[...]
          h:f32[128,128] <- e[...]
          i:f32[128,128] = mul g h
          j:f32[128,128], f[...] <- f[...], i
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[4096,4096]),)
    ] a b
  in (c,) }
Kernel: rsqrt (Level 1, category: elementwise)
Input shapes: [[4096, 4096]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 149.9465 ms; JAX XLA compiled runtime: 0.2054 ms; speedup vs compiled baseline: 0.0014

JAX reference implementation:
@jax.jit
def jax_rsqrt(x: jax.Array) -> jax.Array:
    return jax.lax.rsqrt(x + 1e-5)

Jaxpr IR: { lambda ; a:f32[4096,4096]. let
    b:f32[4096,4096] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(32, 32), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128))), BlockMapping(block_shape=(Blocked(block_size=128), Blocked(block_size=128)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; c:Ref<default>{f32[128,128]} d:Ref<default>{f32[128,128]}. let
          e:f32[128,128] <- c[...]
          f:f32[128,128] = add e 9.999999747378752e-06:f32[]
          g:f32[128,128] = rsqrt f
          h:f32[128,128], d[...] <- d[...], g
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[4096,4096]),)
    ] a
  in (b,) }
Kernel: embedding_lookup (Level 1, category: index)
Input shapes: [[32000, 768], [512]]
Target hardware: NVIDIA A100 80GB PCIe
Status: error (correct: no, max diff vs baseline: -1)
Error: The Pallas Triton lowering currently requires that all operations have array arguments and results whose size is a power of 2. Encountered an array of shape (512, 768)
Pallas runtime: -1 ms; JAX XLA compiled runtime: -1 ms; speedup vs compiled baseline: 0

JAX reference implementation:
@jax.jit
def jax_embedding_lookup(table: jax.Array, indices: jax.Array) -> jax.Array:
    return table[indices]

Original Pallas kernel (TPU-oriented): """Level 1: Embedding table lookup via Pallas.

Demonstrates: gather-style indexing, integer index handling,
non-contiguous memory access patterns.
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _embedding_kernel(table_ref, idx_ref, o_ref):
    idx = idx_ref[...]
    o_ref[...] = table_ref[idx, :]


def pallas_embedding_lookup(
    table: jax.Array, indices: jax.Array
) -> jax.Array:
    seq_len = indices.shape[0]
    embed_dim = table.shape[1]

    return pl.pallas_call(
        _embedding_kernel,
        out_shape=jax.ShapeDtypeStruct((seq_len, embed_dim), table.dtype),
        grid=(1,),
        in_specs=[
            pl.BlockSpec(table.shape, lambda i: (0, 0)),
            pl.BlockSpec(indices.shape, lambda i: (0,)),
        ],
        out_specs=pl.BlockSpec((seq_len, embed_dim), lambda i: (0, 0)),
    )(table, indices)


pallas_kernel = pallas_embedding_lookup
task_name = "embedding_lookup"
input_shapes = [(32000, 768), (512,)]
category = "index"
level = 1

Fixed Pallas kernel (GPU-compatible): """Level 1: Embedding table lookup via Pallas.

Demonstrates: gather-style indexing, integer index handling,
non-contiguous memory access patterns.
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _embedding_kernel(table_ref, idx_ref, o_ref):
    idx = idx_ref[...]
    o_ref[...] = table_ref[idx, :]


def pallas_embedding_lookup(
    table: jax.Array, indices: jax.Array
) -> jax.Array:
    seq_len = indices.shape[0]
    embed_dim = table.shape[1]

    return pl.pallas_call(
        _embedding_kernel,
        out_shape=jax.ShapeDtypeStruct((seq_len, embed_dim), table.dtype),
        grid=(1,),
        in_specs=[
            pl.BlockSpec(table.shape, lambda i: (0, 0)),
            pl.BlockSpec(indices.shape, lambda i: (0,)),
        ],
        out_specs=pl.BlockSpec((seq_len, embed_dim), lambda i: (0, 0)),
    )(table, indices)


pallas_kernel = pallas_embedding_lookup
task_name = "embedding_lookup"
input_shapes = [(32000, 768), (512,)]
category = "index"
level = 1