Kernel: nucleotide_onehot (Level 1, category: genomics)
Input shapes: [[4096]]
Target hardware: NVIDIA A100 80GB PCIe
Status: pass (correct: yes, max diff vs baseline: 0)
Pallas runtime: 83.01 ms; JAX XLA compiled runtime: 0.0987 ms; speedup vs compiled baseline: 0.0012

JAX reference implementation:
@jax.jit
def jax_nucleotide_onehot(seq: jax.Array) -> jax.Array:
    return jax.nn.one_hot(seq, 4, dtype=jnp.float32)

Original Pallas kernel (TPU-oriented): """Level 1: Nucleotide one-hot encoding via Pallas.

Encodes integer-encoded DNA sequences (A=0, C=1, G=2, T=3) into
4-channel one-hot representation used by genomics models (Enformer, etc).

Provenance: google-deepmind/deepmind-research Enformer
             DNA sequence input encoding (one-hot 4-channel)
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _nucleotide_onehot_kernel(seq_ref, o_ref):
    seq = seq_ref[...]
    num_classes = 4
    o_ref[...] = (jnp.arange(num_classes) == seq[..., None]).astype(jnp.float32)


def pallas_nucleotide_onehot(seq: jax.Array) -> jax.Array:
    seq_len = seq.shape[0]

    return pl.pallas_call(
        _nucleotide_onehot_kernel,
        out_shape=jax.ShapeDtypeStruct((seq_len, 4), jnp.float32),
        grid=(1,),
        in_specs=[pl.BlockSpec((seq_len,), lambda i: (0,))],
        out_specs=pl.BlockSpec((seq_len, 4), lambda i: (0, 0)),
    )(seq)


pallas_kernel = pallas_nucleotide_onehot
task_name = "nucleotide_onehot"
input_shapes = [(4096,)]
category = "genomics"
level = 1

Fixed Pallas kernel (GPU-compatible): """Level 1: Nucleotide one-hot encoding via Pallas.

Encodes integer-encoded DNA sequences (A=0, C=1, G=2, T=3) into
4-channel one-hot representation used by genomics models (Enformer, etc).

Provenance: google-deepmind/deepmind-research Enformer
             DNA sequence input encoding (one-hot 4-channel)
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _nucleotide_onehot_kernel(seq_ref, o_ref):
    seq = seq_ref[...]
    num_classes = 4
    o_ref[...] = (jnp.arange(num_classes) == seq[..., None]).astype(jnp.float32)


def pallas_nucleotide_onehot(seq: jax.Array) -> jax.Array:
    seq_len = seq.shape[0]

    return pl.pallas_call(
        _nucleotide_onehot_kernel,
        out_shape=jax.ShapeDtypeStruct((seq_len, 4), jnp.float32),
        grid=(1,),
        in_specs=[pl.BlockSpec((seq_len,), lambda i: (0,))],
        out_specs=pl.BlockSpec((seq_len, 4), lambda i: (0, 0)),
    )(seq)


pallas_kernel = pallas_nucleotide_onehot
task_name = "nucleotide_onehot"
input_shapes = [(4096,)]
category = "genomics"
level = 1

Jaxpr IR: { lambda ; a:i32[4096]. let
    b:f32[4096,4] = pallas_call[
      compiler_params=None
      cost_estimate=None
      debug=False
      grid_mapping=GridMapping(grid=(1,), block_mappings=(BlockMapping(block_shape=(Blocked(block_size=4096),)), BlockMapping(block_shape=(Blocked(block_size=4096), Blocked(block_size=4)))))
      input_output_aliases=()
      interpret=False
      jaxpr={ lambda ; c:Ref<default>{i32[4096]} d:Ref<default>{f32[4096,4]}. let
          e:i32[4096] <- c[...]
          f:i32[4] = iota[dimension=0 dtype=int32 shape=(4,) sharding=None] 
          g:i32[4096,1] = broadcast_in_dim[broadcast_dimensions=(0,)] e
          h:i32[1,4] = broadcast_in_dim[broadcast_dimensions=(1,)] f
          i:bool[4096,4] = eq h g
          j:f32[4096,4] = convert_element_type[
            new_dtype=float32
            weak_type=False
          ] i
          k:f32[4096,4], d[...] <- d[...], j
        in () }
      mesh=None
      metadata=None
      name=None
      out_avals=(ShapedArray(float32[4096,4]),)
    ] a
  in (b,) }