Kernel: sigmoid_bce (Level 2, category: loss_fusion)
Input shapes: [[2048, 1024], [2048, 1024]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 903.8407 ms; JAX XLA compiled runtime: 0.1157 ms; speedup vs compiled baseline: 0.0001

JAX reference implementation:
@jax.jit
def jax_sigmoid_bce(logits: jax.Array, targets: jax.Array) -> jax.Array:
    max_val = jnp.maximum(-logits, 0.0)
    loss = max_val + jnp.log(jnp.exp(-max_val) + jnp.exp(-logits - max_val))
    return loss - targets * logits + targets * loss

Original Pallas kernel (TPU-oriented): """Level 2: Fused Sigmoid + Binary Cross-Entropy via Pallas.

Provenance: standard binary classification loss fusion
"""


from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L2/sigmoid_bce", __doc__)

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


def _sigmoid_bce_kernel(logits_ref, targets_ref, o_ref):
    logits = logits_ref[...]
    targets = targets_ref[...]
    max_val = jnp.maximum(-logits, 0.0)
    loss = max_val + jnp.log(jnp.exp(-max_val) + jnp.exp(-logits - max_val))
    o_ref[...] = loss - targets * logits + targets * loss


def pallas_sigmoid_bce(logits: jax.Array, targets: jax.Array) -> jax.Array:
    n = logits.shape[0]
    block_size = min(1024, n)
    grid_size = n // block_size

    return pl.pallas_call(
        _sigmoid_bce_kernel,
        out_shape=jax.ShapeDtypeStruct(logits.shape, logits.dtype),
        grid=(grid_size,),
        in_specs=[
            pl.BlockSpec((block_size, *logits.shape[1:]), lambda i: (i, *([0] * (logits.ndim - 1)))),
            pl.BlockSpec((block_size, *logits.shape[1:]), lambda i: (i, *([0] * (logits.ndim - 1)))),
        ],
        out_specs=pl.BlockSpec((block_size, *logits.shape[1:]), lambda i: (i, *([0] * (logits.ndim - 1)))),
    )(logits, targets)


pallas_kernel = pallas_sigmoid_bce
task_name = "sigmoid_bce"
input_shapes = [(2048, 1024), (2048, 1024)]
category = "loss_fusion"
level = 2

Fixed Pallas kernel (GPU-compatible): """Level 2: Fused Sigmoid + Binary Cross-Entropy via Pallas.

Provenance: standard binary classification loss fusion
"""


from pallasbench.provenance import describe_task as _describe_task

__doc__ = _describe_task("L2/sigmoid_bce", __doc__)

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


def _sigmoid_bce_kernel(logits_ref, targets_ref, o_ref):
    logits = logits_ref[...]
    targets = targets_ref[...]
    max_val = jnp.maximum(-logits, 0.0)
    loss = max_val + jnp.log(jnp.exp(-max_val) + jnp.exp(-logits - max_val))
    o_ref[...] = loss - targets * logits + targets * loss


def pallas_sigmoid_bce(logits: jax.Array, targets: jax.Array) -> jax.Array:
    n = logits.shape[0]
    cols = 1
    for s in logits.shape[1:]:
        cols *= s
    block_size = min(min(128, n), max(1, 16384 // cols))
    grid_size = n // block_size

    return pl.pallas_call(
        _sigmoid_bce_kernel,
        out_shape=jax.ShapeDtypeStruct(logits.shape, logits.dtype),
        grid=(grid_size,),
        in_specs=[
            pl.BlockSpec((block_size, *logits.shape[1:]), lambda i: (i, *([0] * (logits.ndim - 1)))),
            pl.BlockSpec((block_size, *logits.shape[1:]), lambda i: (i, *([0] * (logits.ndim - 1)))),
        ],
        out_specs=pl.BlockSpec((block_size, *logits.shape[1:]), lambda i: (i, *([0] * (logits.ndim - 1)))),
    )(logits, targets)


pallas_kernel = pallas_sigmoid_bce
task_name = "sigmoid_bce"
input_shapes = [(2048, 1024), (2048, 1024)]
category = "loss_fusion"
level = 2

Diff of GPU fix: --- a/sigmoid_bce.py
+++ b/sigmoid_bce.py
@@ -23,7 +23,10 @@
 
 def pallas_sigmoid_bce(logits: jax.Array, targets: jax.Array) -> jax.Array:
     n = logits.shape[0]
-    block_size = min(1024, n)
+    cols = 1
+    for s in logits.shape[1:]:
+        cols *= s
+    block_size = min(min(128, n), max(1, 16384 // cols))
     grid_size = n // block_size
 
     return pl.pallas_call(

Jaxpr IR: { lambda ; a:f32[2048,1024] b:f32[2048,1024]. let
    c:f32[2048,1024] = 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), Blocked(block_size=1024)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; d:Ref<default>{f32[16,1024]} e:Ref<default>{f32[16,1024]}
          f:Ref<default>{f32[16,1024]}. let
          g:f32[16,1024] <- d[...]
          h:f32[16,1024] <- e[...]
          i:f32[16,1024] = neg g
          j:f32[16,1024] = max i 0.0:f32[]
          k:f32[16,1024] = neg j
          l:f32[16,1024] = exp k
          m:f32[16,1024] = neg g
          n:f32[16,1024] = sub m j
          o:f32[16,1024] = exp n
          p:f32[16,1024] = add l o
          q:f32[16,1024] = log p
          r:f32[16,1024] = add j q
          s:f32[16,1024] = mul h g
          t:f32[16,1024] = sub r s
          u:f32[16,1024] = mul h r
          v:f32[16,1024] = add t u
          w:f32[16,1024], f[...] <- f[...], v
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[2048,1024]),)
    ] a b
  in (c,) }