Kernel: log_softmax (Level 1, category: softmax)
Input shapes: [[2048, 2048]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 1e-06)
Pallas runtime: 312.4499 ms; JAX XLA compiled runtime: 0.1174 ms; speedup vs compiled baseline: 0.0004

JAX reference implementation:
@jax.jit
def jax_log_softmax(x: jax.Array) -> jax.Array:
    return jax.nn.log_softmax(x, axis=-1)

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

Provenance: jax.nn.log_softmax, critical for cross-entropy loss computation
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _log_softmax_kernel(x_ref, o_ref):
    x = x_ref[...]
    row_max = jnp.max(x, axis=-1, keepdims=True)
    shifted = x - row_max
    log_sum_exp = jnp.log(jnp.sum(jnp.exp(shifted), axis=-1, keepdims=True))
    o_ref[...] = shifted - log_sum_exp


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

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


pallas_kernel = pallas_log_softmax
task_name = "log_softmax"
input_shapes = [(2048, 2048)]
category = "softmax"
level = 1

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

Provenance: jax.nn.log_softmax, critical for cross-entropy loss computation
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _log_softmax_kernel(x_ref, o_ref):
    x = x_ref[...]
    row_max = jnp.max(x, axis=-1, keepdims=True)
    shifted = x - row_max
    log_sum_exp = jnp.log(jnp.sum(jnp.exp(shifted), axis=-1, keepdims=True))
    o_ref[...] = shifted - log_sum_exp


def pallas_log_softmax(x: jax.Array) -> jax.Array:
    n_rows = x.shape[0]
    n_cols = x.shape[1]
    block_rows = min(128, 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(
        _log_softmax_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid=(grid_size,),
        in_specs=[pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0))],
        out_specs=pl.BlockSpec((block_rows, n_cols), lambda i: (i, 0)),
    )(x)


pallas_kernel = pallas_log_softmax
task_name = "log_softmax"
input_shapes = [(2048, 2048)]
category = "softmax"
level = 1

Diff of GPU fix: --- a/log_softmax.py
+++ b/log_softmax.py
@@ -25,6 +25,8 @@
     n_rows = x.shape[0]
     n_cols = x.shape[1]
     block_rows = min(128, 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,2048]. let
    b:f32[2048,2048] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(256,), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=8), Blocked(block_size=2048))), BlockMapping(block_shape=(Blocked(block_size=8), Blocked(block_size=2048)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; c:Ref<default>{f32[8,2048]} d:Ref<default>{f32[8,2048]}. let
          e:f32[8,2048] <- c[...]
          f:f32[8] = reduce_max[axes=(1,)] e
          g:f32[8,1] = broadcast_in_dim[broadcast_dimensions=(0,)] f
          h:f32[8,2048] = sub e g
          i:f32[8,2048] = exp h
          j:f32[8] = reduce_sum[axes=(1,) out_sharding=None] i
          k:f32[8,1] = broadcast_in_dim[broadcast_dimensions=(0,)] j
          l:f32[8,1] = log k
          m:f32[8,2048] = sub h l
          n:f32[8,2048], d[...] <- d[...], m
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[2048,2048]),)
    ] a
  in (b,) }