[Pallas TPU] Incorrect result for jnp.dot when rhs is uint4 and lhs is bfloat16

[catswe]

Description

The two Pallas calls below should have produced equivalent results. I would have expected an unimplemented error raised if this is truly not supported

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

def dot(lhs_ref, rhs_ref, out_ref):
    out_ref[...] = jnp.dot(lhs_ref[...], rhs_ref[...], preferred_element_type=jnp.float32)

M, K, N = 128, 256, 128
lhs = jax.random.normal(jax.random.PRNGKey(0), (M, K), dtype=jnp.bfloat16)
rhs = jax.random.randint(jax.random.PRNGKey(1), (K, N), 0, 16).astype(jnp.uint4)

vmem = pl.BlockSpec(memory_space=pltpu.VMEM)
call_args = dict(
    out_shape=jax.ShapeDtypeStruct((M, N), jnp.float32),
    in_specs=[vmem, vmem],
    out_specs=vmem,
)

result_uncasted = pl.pallas_call(dot, **call_args)(lhs, rhs) # rhs is uint4
result_casted  = pl.pallas_call(dot, **call_args)(lhs, rhs.astype(jnp.bfloat16))

diff = jnp.abs(result_uncasted - result_casted)
print(f"uncasted result sample: {result_uncasted[0, :4]}")
print(f"casted result sample: {result_casted[0, :4]}")
print(f"max diff:  {jnp.max(diff)}")
print(f"mean diff: {jnp.mean(diff)}")
print(f"allclose (atol=1e-3): {jnp.allclose(result_uncasted, result_casted, atol=1e-3)}")

uncasted result sample: [ 32.20056    18.895325  -71.4046      7.1782227]
casted result sample: [ 38.04431    50.196106    6.5143433 -61.946777 ]
max diff:  875.46484375
mean diff: 142.3582305908203
allclose (atol=1e-3): False

System info (python version, jaxlib version, accelerator, etc.)

>>> import jax; jax.print_environment_info()
jax:    0.8.1
jaxlib: 0.8.1
numpy:  2.3.5
python: 3.11.13 (main, Jun  4 2025, 08:57:29) [GCC 11.4.0]
device info: TPU v6 lite-8, 8 local devices"
process_count: 1
platform: uname_result(system='Linux', node='t1v-n-f16f40fa-w-0', release='6.8.0-1015-gcp', version='#17~22.04.1-Ubuntu SMP Tue Sep  3 16:11:52 UTC 2024', machine='x86_64')

[rdyro]

This looks like a problem with uint4 specifically, since when I change

- rhs = jax.random.randint(jax.random.PRNGKey(1), (K, N), 0, 16).astype(jnp.uint4)
+ rhs = jax.random.randint(jax.random.PRNGKey(1), (K, N), -8, 7).astype(jnp.int4)

your test passes. Let me look into this.

[catswe]

Closing since there's a fix commit