Kernel: pairwise_distance (Level 2, category: genomics)
Input shapes: [[256, 32]]
Target hardware: NVIDIA A100 80GB PCIe
Status: error (correct: no, max diff vs baseline: -1)
Error: Verification failed:
error: "reduce_sum:"("reduce_sum"("_pairwise_dist_kernel"("/mnt/data2/pallasbench_fix/pallasbench/kernels/level2/pairwise_distance.py":25:26 to :55))): Maximum allowed number of e
Pallas runtime: -1 ms; JAX XLA compiled runtime: -1 ms; speedup vs compiled baseline: 0

JAX reference implementation:
@jax.jit
def jax_pairwise_distance(x: jax.Array) -> jax.Array:
    diff = x[:, None, :] - x[None, :, :]
    return jnp.sqrt(jnp.sum(diff * diff, axis=-1) + 1e-8)

Original Pallas kernel (TPU-oriented): """Level 2: Pairwise Euclidean distance matrix via Pallas.

Computes the N x N distance matrix from N points in d dimensions.
Core operation in structural biology (AlphaFold distance maps),
molecular dynamics (neighbor lists), and genomics (phylogenetics).

Provenance: google-deepmind/alphafold3 pair representation distance maps
             JAX-MD (arxiv:1912.04232) pairwise distance computation
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _pairwise_dist_kernel(x_ref, o_ref):
    x = x_ref[...]
    n = x.shape[0]
    diff = x[:, None, :] - x[None, :, :]
    o_ref[...] = jnp.sqrt(jnp.sum(diff * diff, axis=-1) + 1e-8)


def pallas_pairwise_distance(x: jax.Array) -> jax.Array:
    n, d = x.shape

    return pl.pallas_call(
        _pairwise_dist_kernel,
        out_shape=jax.ShapeDtypeStruct((n, n), x.dtype),
        grid=(1,),
        in_specs=[pl.BlockSpec((n, d), lambda i: (0, 0))],
        out_specs=pl.BlockSpec((n, n), lambda i: (0, 0)),
    )(x)


pallas_kernel = pallas_pairwise_distance
task_name = "pairwise_distance"
input_shapes = [(256, 32)]
category = "genomics"
level = 2

Fixed Pallas kernel (GPU-compatible): """Level 2: Pairwise Euclidean distance matrix via Pallas.

Computes the N x N distance matrix from N points in d dimensions.
Core operation in structural biology (AlphaFold distance maps),
molecular dynamics (neighbor lists), and genomics (phylogenetics).

Provenance: google-deepmind/alphafold3 pair representation distance maps
             JAX-MD (arxiv:1912.04232) pairwise distance computation
"""


from pallasbench.provenance import describe_task as _describe_task

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

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


def _pairwise_dist_kernel(x_ref, o_ref):
    x = x_ref[...]
    n = x.shape[0]
    diff = x[:, None, :] - x[None, :, :]
    o_ref[...] = jnp.sqrt(jnp.sum(diff * diff, axis=-1) + 1e-8)


def pallas_pairwise_distance(x: jax.Array) -> jax.Array:
    n, d = x.shape

    return pl.pallas_call(
        _pairwise_dist_kernel,
        out_shape=jax.ShapeDtypeStruct((n, n), x.dtype),
        grid=(1,),
        in_specs=[pl.BlockSpec((n, d), lambda i: (0, 0))],
        out_specs=pl.BlockSpec((n, n), lambda i: (0, 0)),
    )(x)


pallas_kernel = pallas_pairwise_distance
task_name = "pairwise_distance"
input_shapes = [(256, 32)]
category = "genomics"
level = 2