Silent wrong results for an integer-only StableHLO program from jax.random (llvm-cpu and vmvx agree with each other, disagree with XLA)

[maraxen]

What happened?

A small JAX program lowered to StableHLO returns different numerical results under IREE than under XLA, silently — no error, no warning, and the values look entirely plausible.

The program is integer arithmetic only (xor / shift_left / shift_right_logical / add inside two stablehlo.while loops). There is no stablehlo.rng_bit_generator, no custom call, and no floating-point reassociation available as an explanation — the wrong values are exactly reproducible and identical run to run.

It is Threefry, as emitted by jax.random, so in practice this silently changes every sampled value in a model exported through IREE.

Expected (eager JAX, and XLA running the very same exported StableHLO module — see "whose bug is it" below):

result[0] 2x8xui32 = [1214273199 3384852239 1707608394 3966931780 1289383684 173908716 827220724 2504598792]
result[1] 2x8xf32  = [-2.4424558 -2.0356805 0.20554423 -0.3535502 -0.76197404 -1.1785518 -1.1482196 0.29716578]

Actual (IREE):

result[0] 2x8xui32 = [3223025141 1838280444 2167582192 563678655 2907208703 835636454 3197201895 1574853171]
result[1] 2x8xf32  = [-0.077075 0.994366 0.202603 1.84414 0.311271 1.64625 -0.853936 -1.2273]

Every element of both outputs is wrong. Both lanes are identical in both cases (the two lanes hold the same key), so this is not lane-dependent.

Whose bug is it?

I checked this before filing rather than assuming. jax.export's Exported.call re-imports the same StableHLO module and runs it through XLA:

	vs eager JAX
XLA, running the exported module	0 / 16 elements differ
IREE, running the exported module	16 / 16 elements differ

So the StableHLO being handed to IREE is correct, and IREE miscompiles it.

Minimization

Reduced from a real model (a ProteinMPNN-family scorer) down to this, by removal, with a live control at every step. Two ingredients are each necessary — remove either one and IREE is bit-exact:

1. a nested jax.random.split — a split of a split half — feeding the draw, and
2. a second live random draw elsewhere in the program.

Notably, a single-level split with both halves drawn from is exact. So this is not simply "two draws"; something about the extra derivation depth matters. I have not isolated which StableHLO op is at fault — the minimization above is at the JAX level, and the attached MLIR is the artifact.

Removing jax.vmap does not make it exact — it makes the runtime abort instead (see Additional context).

Steps to reproduce your issue

The attached MLIR reproduces with the CLI alone, no JAX needed:

$ iree-compile repro.mlir --iree-hal-target-backends=llvm-cpu \
    --iree-llvmcpu-target-cpu=host -o repro.vmfb
$ iree-run-module --module=repro.vmfb --function=main --input=2x2xi32=0,0,0,0
EXEC @main
result[0]: hal.buffer_view
2x8xui32=[3223025141 1838280444 ...]     # wrong; see Expected above

To regenerate the MLIR from source (jax 0.11.1):

import jax, jax.numpy as jnp
from jax import export as jexport

def minimal(key):
    k1, k2 = jax.random.split(key)
    nested, _ = jax.random.split(k1)          # the nested split
    return (jax.random.bits(nested, (8,), dtype=jnp.uint32),
            jax.random.normal(k2, (8,), dtype=jnp.float32))   # the second draw

keys = jnp.tile(jnp.array([0, 0], dtype=jnp.uint32), (2, 1))
exported = jexport.export(jax.jit(jax.vmap(minimal)))(
    jax.ShapeDtypeStruct(keys.shape, keys.dtype))
open("repro.mlir", "w").write(exported.mlir_module())

print(exported.call(keys))   # XLA on this exact module -- correct

repro.mlir (437 lines)

#loc = loc(unknown)
#loc1 = loc("key")
#loc2 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5093/prefile.py":37:4 to :51)
#loc5 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5093/prefile.py":34:11 to :32)
#loc23 = loc("minimal"(#loc2))
#loc26 = loc("minimal"(#loc5))
#loc48 = loc("vmap()"(#loc23))
module @jit_minimal attributes {jax.uses_shape_polymorphism = false, mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<2x2xui32> loc("key")) -> (tensor<2x8xui32> {jax.result_info = "result[0]"}, tensor<2x8xf32> {jax.result_info = "result[1]"}) {
    %c = stablehlo.constant dense<32> : tensor<ui64> loc(#loc90)
    %c_0 = stablehlo.constant dense<1> : tensor<ui64> loc(#loc90)
    %0 = call @_threefry_split(%arg0) : (tensor<2x2xui32>) -> tensor<2x2x2xui32> loc(#loc91)
    %1 = stablehlo.slice %0 [0:2, 0:1, 0:2] : (tensor<2x2x2xui32>) -> tensor<2x1x2xui32> loc(#loc79)
    %2 = stablehlo.reshape %1 : (tensor<2x1x2xui32>) -> tensor<2x2xui32> loc(#loc79)
    %3 = stablehlo.slice %0 [0:2, 1:2, 0:2] : (tensor<2x2x2xui32>) -> tensor<2x1x2xui32> loc(#loc79)
    %4 = stablehlo.reshape %3 : (tensor<2x1x2xui32>) -> tensor<2x2xui32> loc(#loc79)
    %5 = call @_threefry_split(%2) : (tensor<2x2xui32>) -> tensor<2x2x2xui32> loc(#loc92)
    %6 = stablehlo.slice %5 [0:2, 0:1, 0:2] : (tensor<2x2x2xui32>) -> tensor<2x1x2xui32> loc(#loc81)
    %7 = stablehlo.reshape %6 : (tensor<2x1x2xui32>) -> tensor<2x2xui32> loc(#loc81)
    %8 = stablehlo.slice %7 [0:2, 0:1] : (tensor<2x2xui32>) -> tensor<2x1xui32> loc(#loc93)
    %9 = stablehlo.reshape %8 : (tensor<2x1xui32>) -> tensor<2xui32> loc(#loc93)
    %10 = stablehlo.slice %7 [0:2, 1:2] : (tensor<2x2xui32>) -> tensor<2x1xui32> loc(#loc93)
    %11 = stablehlo.reshape %10 : (tensor<2x1xui32>) -> tensor<2xui32> loc(#loc93)
    %12 = stablehlo.iota dim = 0 : tensor<8xui64> loc(#loc90)
    %13 = stablehlo.broadcast_in_dim %c_0, dims = [] : (tensor<ui64>) -> tensor<8xui64> loc(#loc90)
    %14 = stablehlo.multiply %13, %12 : tensor<8xui64> loc(#loc90)
    %15 = stablehlo.broadcast_in_dim %c, dims = [] : (tensor<ui64>) -> tensor<8xui64> loc(#loc90)
    %16 = stablehlo.shift_right_logical %14, %15 : tensor<8xui64> loc(#loc90)
    %17 = stablehlo.convert %14 : (tensor<8xui64>) -> tensor<8xui32> loc(#loc90)
    %18 = stablehlo.convert %16 : (tensor<8xui64>) -> tensor<8xui32> loc(#loc90)
    %19 = stablehlo.broadcast_in_dim %18, dims = [1] : (tensor<8xui32>) -> tensor<1x8xui32> loc(#loc94)
    %20 = stablehlo.broadcast_in_dim %17, dims = [1] : (tensor<8xui32>) -> tensor<1x8xui32> loc(#loc94)
    %21 = stablehlo.broadcast_in_dim %9, dims = [0] : (tensor<2xui32>) -> tensor<2x1xui32> loc(#loc94)
    %22 = stablehlo.broadcast_in_dim %11, dims = [0] : (tensor<2xui32>) -> tensor<2x1xui32> loc(#loc94)
    %23:2 = call @threefry2x32_0(%21, %22, %19, %20) : (tensor<2x1xui32>, tensor<2x1xui32>, tensor<1x8xui32>, tensor<1x8xui32>) -> (tensor<2x8xui32>, tensor<2x8xui32>) loc(#loc95)
    %24 = stablehlo.xor %23#0, %23#1 : tensor<2x8xui32> loc(#loc96)
    %25 = call @_normal(%4) : (tensor<2x2xui32>) -> tensor<2x8xf32> loc(#loc82)
    return %24, %25 : tensor<2x8xui32>, tensor<2x8xf32> loc(#loc69)
  } loc(#loc)
  func.func private @_threefry_split(%arg0: tensor<2x2xui32> loc(unknown)) -> tensor<2x2x2xui32> {
    %c = stablehlo.constant dense<32> : tensor<ui64> loc(#loc35)
    %c_0 = stablehlo.constant dense<1> : tensor<ui64> loc(#loc35)
    %0 = stablehlo.slice %arg0 [0:2, 0:1] : (tensor<2x2xui32>) -> tensor<2x1xui32> loc(#loc36)
    %1 = stablehlo.reshape %0 : (tensor<2x1xui32>) -> tensor<2xui32> loc(#loc36)
    %2 = stablehlo.slice %arg0 [0:2, 1:2] : (tensor<2x2xui32>) -> tensor<2x1xui32> loc(#loc36)
    %3 = stablehlo.reshape %2 : (tensor<2x1xui32>) -> tensor<2xui32> loc(#loc36)
    %4 = stablehlo.iota dim = 0 : tensor<2xui64> loc(#loc35)
    %5 = stablehlo.broadcast_in_dim %c_0, dims = [] : (tensor<ui64>) -> tensor<2xui64> loc(#loc35)
    %6 = stablehlo.multiply %5, %4 : tensor<2xui64> loc(#loc35)
    %7 = stablehlo.broadcast_in_dim %c, dims = [] : (tensor<ui64>) -> tensor<2xui64> loc(#loc35)
    %8 = stablehlo.shift_right_logical %6, %7 : tensor<2xui64> loc(#loc35)
    %9 = stablehlo.convert %6 : (tensor<2xui64>) -> tensor<2xui32> loc(#loc35)
    %10 = stablehlo.convert %8 : (tensor<2xui64>) -> tensor<2xui32> loc(#loc35)
    %11 = stablehlo.broadcast_in_dim %10, dims = [1] : (tensor<2xui32>) -> tensor<1x2xui32> loc(#loc37)
    %12 = stablehlo.broadcast_in_dim %9, dims = [1] : (tensor<2xui32>) -> tensor<1x2xui32> loc(#loc37)
    %13 = stablehlo.broadcast_in_dim %1, dims = [0] : (tensor<2xui32>) -> tensor<2x1xui32> loc(#loc37)
    %14 = stablehlo.broadcast_in_dim %3, dims = [0] : (tensor<2xui32>) -> tensor<2x1xui32> loc(#loc37)
    %15:2 = call @threefry2x32(%13, %14, %11, %12) : (tensor<2x1xui32>, tensor<2x1xui32>, tensor<1x2xui32>, tensor<1x2xui32>) -> (tensor<2x2xui32>, tensor<2x2xui32>) loc(#loc26)
    %16 = stablehlo.broadcast_in_dim %15#0, dims = [0, 1] : (tensor<2x2xui32>) -> tensor<2x2x1xui32> loc(#loc38)
    %17 = stablehlo.broadcast_in_dim %15#1, dims = [0, 1] : (tensor<2x2xui32>) -> tensor<2x2x1xui32> loc(#loc38)
    %18 = stablehlo.concatenate %16, %17, dim = 2 : (tensor<2x2x1xui32>, tensor<2x2x1xui32>) -> tensor<2x2x2xui32> loc(#loc38)
    return %18 : tensor<2x2x2xui32> loc(#loc)
  } loc(#loc70)
  func.func private @threefry2x32(%arg0: tensor<2x1xui32> loc("minimal"(#loc5)), %arg1: tensor<2x1xui32> loc("minimal"(#loc5)), %arg2: tensor<1x2xui32> loc("minimal"(#loc5)), %arg3: tensor<1x2xui32> loc("minimal"(#loc5))) -> (tensor<2x2xui32>, tensor<2x2xui32>) {
    %c = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %c_0 = stablehlo.constant dense<5> : tensor<i32> loc(#loc)
    %c_1 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %c_2 = stablehlo.constant dense<466688986> : tensor<ui32> loc(#loc31)
    %c_3 = stablehlo.constant dense<[13, 15, 26, 6]> : tensor<4xui32> loc(#loc31)
    %c_4 = stablehlo.constant dense<[17, 29, 16, 24]> : tensor<4xui32> loc(#loc31)
    %0 = stablehlo.xor %arg0, %arg1 : tensor<2x1xui32> loc(#loc39)
    %1 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<ui32>) -> tensor<2x1xui32> loc(#loc39)
    %2 = stablehlo.xor %0, %1 : tensor<2x1xui32> loc(#loc39)
    %3 = stablehlo.broadcast_in_dim %arg2, dims = [0, 1] : (tensor<1x2xui32>) -> tensor<2x2xui32> loc(#loc40)
    %4 = stablehlo.broadcast_in_dim %arg0, dims = [0, 1] : (tensor<2x1xui32>) -> tensor<2x2xui32> loc(#loc40)
    %5 = stablehlo.add %3, %4 : tensor<2x2xui32> loc(#loc40)
    %6 = stablehlo.broadcast_in_dim %arg3, dims = [0, 1] : (tensor<1x2xui32>) -> tensor<2x2xui32> loc(#loc40)
    %7 = stablehlo.broadcast_in_dim %arg1, dims = [0, 1] : (tensor<2x1xui32>) -> tensor<2x2xui32> loc(#loc40)
    %8 = stablehlo.add %6, %7 : tensor<2x2xui32> loc(#loc40)
    %9:9 = stablehlo.while(%iterArg = %c_1, %iterArg_5 = %c_1, %iterArg_6 = %5, %iterArg_7 = %8, %iterArg_8 = %arg1, %iterArg_9 = %2, %iterArg_10 = %arg0, %iterArg_11 = %c_3, %iterArg_12 = %c_4) : tensor<i32>, tensor<i32>, tensor<2x2xui32>, tensor<2x2xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32>
    cond {
      %10 = stablehlo.compare LT, %iterArg, %c_0, SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc83)
      stablehlo.return %10 : tensor<i1> loc(#loc71)
    } do {
      %10:8 = func.call @eval_jaxpr(%iterArg_5, %iterArg_6, %iterArg_7, %iterArg_8, %iterArg_9, %iterArg_10, %iterArg_11, %iterArg_12) : (tensor<i32>, tensor<2x2xui32>, tensor<2x2xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32>) -> (tensor<i32>, tensor<2x2xui32>, tensor<2x2xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32>) loc(#loc84)
      %11 = stablehlo.add %iterArg, %c : tensor<i32> loc(#loc85)
      stablehlo.return %11, %10#0, %10#1, %10#2, %10#3, %10#4, %10#5, %10#6, %10#7 : tensor<i32>, tensor<i32>, tensor<2x2xui32>, tensor<2x2xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32> loc(#loc71)
    } loc(#loc71)
    return %9#2, %9#3 : tensor<2x2xui32>, tensor<2x2xui32> loc(#loc26)
  } loc(#loc26)
  func.func private @eval_jaxpr(%arg0: tensor<i32> loc(unknown), %arg1: tensor<2x2xui32> loc(unknown), %arg2: tensor<2x2xui32> loc(unknown), %arg3: tensor<2x1xui32> loc(unknown), %arg4: tensor<2x1xui32> loc(unknown), %arg5: tensor<2x1xui32> loc(unknown), %arg6: tensor<4xui32> loc(unknown), %arg7: tensor<4xui32> loc(unknown)) -> (tensor<i32>, tensor<2x2xui32>, tensor<2x2xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32>) {
    %c = stablehlo.constant dense<32> : tensor<ui32> loc(#loc)
    %c_0 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %0 = stablehlo.add %arg0, %c_0 : tensor<i32> loc(#loc40)
    %1 = stablehlo.slice %arg6 [0:1] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc36)
    %2 = stablehlo.reshape %1 : (tensor<1xui32>) -> tensor<ui32> loc(#loc36)
    %3 = stablehlo.slice %arg6 [1:2] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc36)
    %4 = stablehlo.reshape %3 : (tensor<1xui32>) -> tensor<ui32> loc(#loc36)
    %5 = stablehlo.slice %arg6 [2:3] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc36)
    %6 = stablehlo.reshape %5 : (tensor<1xui32>) -> tensor<ui32> loc(#loc36)
    %7 = stablehlo.slice %arg6 [3:4] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc36)
    %8 = stablehlo.reshape %7 : (tensor<1xui32>) -> tensor<ui32> loc(#loc36)
    %9 = stablehlo.add %arg1, %arg2 : tensor<2x2xui32> loc(#loc40)
    %10 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<ui32>) -> tensor<2x2xui32> loc(#loc43)
    %11 = stablehlo.shift_left %arg2, %10 : tensor<2x2xui32> loc(#loc43)
    %12 = stablehlo.subtract %c, %2 : tensor<ui32> loc(#loc44)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<ui32>) -> tensor<2x2xui32> loc(#loc45)
    %14 = stablehlo.shift_right_logical %arg2, %13 : tensor<2x2xui32> loc(#loc45)
    %15 = stablehlo.or %11, %14 : tensor<2x2xui32> loc(#loc46)
    %16 = stablehlo.xor %9, %15 : tensor<2x2xui32> loc(#loc39)
    %17 = stablehlo.add %9, %16 : tensor<2x2xui32> loc(#loc40)
    %18 = stablehlo.broadcast_in_dim %4, dims = [] : (tensor<ui32>) -> tensor<2x2xui32> loc(#loc43)
    %19 = stablehlo.shift_left %16, %18 : tensor<2x2xui32> loc(#loc43)
    %20 = stablehlo.subtract %c, %4 : tensor<ui32> loc(#loc44)
    %21 = stablehlo.broadcast_in_dim %20, dims = [] : (tensor<ui32>) -> tensor<2x2xui32> loc(#loc45)
    %22 = stablehlo.shift_right_logical %16, %21 : tensor<2x2xui32> loc(#loc45)
    %23 = stablehlo.or %19, %22 : tensor<2x2xui32> loc(#loc46)
    %24 = stablehlo.xor %17, %23 : tensor<2x2xui32> loc(#loc39)
    %25 = stablehlo.add %17, %24 : tensor<2x2xui32> loc(#loc40)
    %26 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor<ui32>) -> tensor<2x2xui32> loc(#loc43)
    %27 = stablehlo.shift_left %24, %26 : tensor<2x2xui32> loc(#loc43)
    %28 = stablehlo.subtract %c, %6 : tensor<ui32> loc(#loc44)
    %29 = stablehlo.broadcast_in_dim %28, dims = [] : (tensor<ui32>) -> tensor<2x2xui32> loc(#loc45)
    %30 = stablehlo.shift_right_logical %24, %29 : tensor<2x2xui32> loc(#loc45)
    %31 = stablehlo.or %27, %30 : tensor<2x2xui32> loc(#loc46)
    %32 = stablehlo.xor %25, %31 : tensor<2x2xui32> loc(#loc39)
    %33 = stablehlo.add %25, %32 : tensor<2x2xui32> loc(#loc40)
    %34 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<ui32>) -> tensor<2x2xui32> loc(#loc43)
    %35 = stablehlo.shift_left %32, %34 : tensor<2x2xui32> loc(#loc43)
    %36 = stablehlo.subtract %c, %8 : tensor<ui32> loc(#loc44)
    %37 = stablehlo.broadcast_in_dim %36, dims = [] : (tensor<ui32>) -> tensor<2x2xui32> loc(#loc45)
    %38 = stablehlo.shift_right_logical %32, %37 : tensor<2x2xui32> loc(#loc45)
    %39 = stablehlo.or %35, %38 : tensor<2x2xui32> loc(#loc46)
    %40 = stablehlo.xor %33, %39 : tensor<2x2xui32> loc(#loc39)
    %41 = stablehlo.broadcast_in_dim %arg3, dims = [0, 1] : (tensor<2x1xui32>) -> tensor<2x2xui32> loc(#loc40)
    %42 = stablehlo.add %33, %41 : tensor<2x2xui32> loc(#loc40)
    %43 = stablehlo.broadcast_in_dim %arg4, dims = [0, 1] : (tensor<2x1xui32>) -> tensor<2x2xui32> loc(#loc40)
    %44 = stablehlo.add %40, %43 : tensor<2x2xui32> loc(#loc40)
    %45 = stablehlo.add %arg0, %c_0 : tensor<i32> loc(#loc40)
    %46 = stablehlo.convert %45 : (tensor<i32>) -> tensor<ui32> loc(#loc47)
    %47 = stablehlo.broadcast_in_dim %46, dims = [] : (tensor<ui32>) -> tensor<2x2xui32> loc(#loc40)
    %48 = stablehlo.add %44, %47 : tensor<2x2xui32> loc(#loc40)
    return %0, %42, %48, %arg4, %arg5, %arg3, %arg7, %arg6 : tensor<i32>, tensor<2x2xui32>, tensor<2x2xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32> loc(#loc)
  } loc(#loc73)
  func.func private @threefry2x32_0(%arg0: tensor<2x1xui32> loc("vmap()"(#loc23)), %arg1: tensor<2x1xui32> loc("vmap()"(#loc23)), %arg2: tensor<1x8xui32> loc("vmap()"(#loc23)), %arg3: tensor<1x8xui32> loc("vmap()"(#loc23))) -> (tensor<2x8xui32>, tensor<2x8xui32>) {
    %c = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %c_0 = stablehlo.constant dense<5> : tensor<i32> loc(#loc)
    %c_1 = stablehlo.constant dense<0> : tensor<i32> loc(#loc)
    %c_2 = stablehlo.constant dense<466688986> : tensor<ui32> loc(#loc31)
    %c_3 = stablehlo.constant dense<[13, 15, 26, 6]> : tensor<4xui32> loc(#loc31)
    %c_4 = stablehlo.constant dense<[17, 29, 16, 24]> : tensor<4xui32> loc(#loc31)
    %0 = stablehlo.xor %arg0, %arg1 : tensor<2x1xui32> loc(#loc49)
    %1 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<ui32>) -> tensor<2x1xui32> loc(#loc49)
    %2 = stablehlo.xor %0, %1 : tensor<2x1xui32> loc(#loc49)
    %3 = stablehlo.broadcast_in_dim %arg2, dims = [0, 1] : (tensor<1x8xui32>) -> tensor<2x8xui32> loc(#loc50)
    %4 = stablehlo.broadcast_in_dim %arg0, dims = [0, 1] : (tensor<2x1xui32>) -> tensor<2x8xui32> loc(#loc50)
    %5 = stablehlo.add %3, %4 : tensor<2x8xui32> loc(#loc50)
    %6 = stablehlo.broadcast_in_dim %arg3, dims = [0, 1] : (tensor<1x8xui32>) -> tensor<2x8xui32> loc(#loc50)
    %7 = stablehlo.broadcast_in_dim %arg1, dims = [0, 1] : (tensor<2x1xui32>) -> tensor<2x8xui32> loc(#loc50)
    %8 = stablehlo.add %6, %7 : tensor<2x8xui32> loc(#loc50)
    %9:9 = stablehlo.while(%iterArg = %c_1, %iterArg_5 = %c_1, %iterArg_6 = %5, %iterArg_7 = %8, %iterArg_8 = %arg1, %iterArg_9 = %2, %iterArg_10 = %arg0, %iterArg_11 = %c_3, %iterArg_12 = %c_4) : tensor<i32>, tensor<i32>, tensor<2x8xui32>, tensor<2x8xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32>
    cond {
      %10 = stablehlo.compare LT, %iterArg, %c_0, SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc86)
      stablehlo.return %10 : tensor<i1> loc(#loc74)
    } do {
      %10:8 = func.call @eval_jaxpr_1(%iterArg_5, %iterArg_6, %iterArg_7, %iterArg_8, %iterArg_9, %iterArg_10, %iterArg_11, %iterArg_12) : (tensor<i32>, tensor<2x8xui32>, tensor<2x8xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32>) -> (tensor<i32>, tensor<2x8xui32>, tensor<2x8xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32>) loc(#loc87)
      %11 = stablehlo.add %iterArg, %c : tensor<i32> loc(#loc88)
      stablehlo.return %11, %10#0, %10#1, %10#2, %10#3, %10#4, %10#5, %10#6, %10#7 : tensor<i32>, tensor<i32>, tensor<2x8xui32>, tensor<2x8xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32> loc(#loc74)
    } loc(#loc74)
    return %9#2, %9#3 : tensor<2x8xui32>, tensor<2x8xui32> loc(#loc48)
  } loc(#loc48)
  func.func private @eval_jaxpr_1(%arg0: tensor<i32> loc(unknown), %arg1: tensor<2x8xui32> loc(unknown), %arg2: tensor<2x8xui32> loc(unknown), %arg3: tensor<2x1xui32> loc(unknown), %arg4: tensor<2x1xui32> loc(unknown), %arg5: tensor<2x1xui32> loc(unknown), %arg6: tensor<4xui32> loc(unknown), %arg7: tensor<4xui32> loc(unknown)) -> (tensor<i32>, tensor<2x8xui32>, tensor<2x8xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32>) {
    %c = stablehlo.constant dense<32> : tensor<ui32> loc(#loc)
    %c_0 = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %0 = stablehlo.add %arg0, %c_0 : tensor<i32> loc(#loc50)
    %1 = stablehlo.slice %arg6 [0:1] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc53)
    %2 = stablehlo.reshape %1 : (tensor<1xui32>) -> tensor<ui32> loc(#loc53)
    %3 = stablehlo.slice %arg6 [1:2] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc53)
    %4 = stablehlo.reshape %3 : (tensor<1xui32>) -> tensor<ui32> loc(#loc53)
    %5 = stablehlo.slice %arg6 [2:3] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc53)
    %6 = stablehlo.reshape %5 : (tensor<1xui32>) -> tensor<ui32> loc(#loc53)
    %7 = stablehlo.slice %arg6 [3:4] : (tensor<4xui32>) -> tensor<1xui32> loc(#loc53)
    %8 = stablehlo.reshape %7 : (tensor<1xui32>) -> tensor<ui32> loc(#loc53)
    %9 = stablehlo.add %arg1, %arg2 : tensor<2x8xui32> loc(#loc50)
    %10 = stablehlo.broadcast_in_dim %2, dims = [] : (tensor<ui32>) -> tensor<2x8xui32> loc(#loc54)
    %11 = stablehlo.shift_left %arg2, %10 : tensor<2x8xui32> loc(#loc54)
    %12 = stablehlo.subtract %c, %2 : tensor<ui32> loc(#loc55)
    %13 = stablehlo.broadcast_in_dim %12, dims = [] : (tensor<ui32>) -> tensor<2x8xui32> loc(#loc56)
    %14 = stablehlo.shift_right_logical %arg2, %13 : tensor<2x8xui32> loc(#loc56)
    %15 = stablehlo.or %11, %14 : tensor<2x8xui32> loc(#loc57)
    %16 = stablehlo.xor %9, %15 : tensor<2x8xui32> loc(#loc49)
    %17 = stablehlo.add %9, %16 : tensor<2x8xui32> loc(#loc50)
    %18 = stablehlo.broadcast_in_dim %4, dims = [] : (tensor<ui32>) -> tensor<2x8xui32> loc(#loc54)
    %19 = stablehlo.shift_left %16, %18 : tensor<2x8xui32> loc(#loc54)
    %20 = stablehlo.subtract %c, %4 : tensor<ui32> loc(#loc55)
    %21 = stablehlo.broadcast_in_dim %20, dims = [] : (tensor<ui32>) -> tensor<2x8xui32> loc(#loc56)
    %22 = stablehlo.shift_right_logical %16, %21 : tensor<2x8xui32> loc(#loc56)
    %23 = stablehlo.or %19, %22 : tensor<2x8xui32> loc(#loc57)
    %24 = stablehlo.xor %17, %23 : tensor<2x8xui32> loc(#loc49)
    %25 = stablehlo.add %17, %24 : tensor<2x8xui32> loc(#loc50)
    %26 = stablehlo.broadcast_in_dim %6, dims = [] : (tensor<ui32>) -> tensor<2x8xui32> loc(#loc54)
    %27 = stablehlo.shift_left %24, %26 : tensor<2x8xui32> loc(#loc54)
    %28 = stablehlo.subtract %c, %6 : tensor<ui32> loc(#loc55)
    %29 = stablehlo.broadcast_in_dim %28, dims = [] : (tensor<ui32>) -> tensor<2x8xui32> loc(#loc56)
    %30 = stablehlo.shift_right_logical %24, %29 : tensor<2x8xui32> loc(#loc56)
    %31 = stablehlo.or %27, %30 : tensor<2x8xui32> loc(#loc57)
    %32 = stablehlo.xor %25, %31 : tensor<2x8xui32> loc(#loc49)
    %33 = stablehlo.add %25, %32 : tensor<2x8xui32> loc(#loc50)
    %34 = stablehlo.broadcast_in_dim %8, dims = [] : (tensor<ui32>) -> tensor<2x8xui32> loc(#loc54)
    %35 = stablehlo.shift_left %32, %34 : tensor<2x8xui32> loc(#loc54)
    %36 = stablehlo.subtract %c, %8 : tensor<ui32> loc(#loc55)
    %37 = stablehlo.broadcast_in_dim %36, dims = [] : (tensor<ui32>) -> tensor<2x8xui32> loc(#loc56)
    %38 = stablehlo.shift_right_logical %32, %37 : tensor<2x8xui32> loc(#loc56)
    %39 = stablehlo.or %35, %38 : tensor<2x8xui32> loc(#loc57)
    %40 = stablehlo.xor %33, %39 : tensor<2x8xui32> loc(#loc49)
    %41 = stablehlo.broadcast_in_dim %arg3, dims = [0, 1] : (tensor<2x1xui32>) -> tensor<2x8xui32> loc(#loc50)
    %42 = stablehlo.add %33, %41 : tensor<2x8xui32> loc(#loc50)
    %43 = stablehlo.broadcast_in_dim %arg4, dims = [0, 1] : (tensor<2x1xui32>) -> tensor<2x8xui32> loc(#loc50)
    %44 = stablehlo.add %40, %43 : tensor<2x8xui32> loc(#loc50)
    %45 = stablehlo.add %arg0, %c_0 : tensor<i32> loc(#loc50)
    %46 = stablehlo.convert %45 : (tensor<i32>) -> tensor<ui32> loc(#loc58)
    %47 = stablehlo.broadcast_in_dim %46, dims = [] : (tensor<ui32>) -> tensor<2x8xui32> loc(#loc50)
    %48 = stablehlo.add %44, %47 : tensor<2x8xui32> loc(#loc50)
    return %0, %42, %48, %arg4, %arg5, %arg3, %arg7, %arg6 : tensor<i32>, tensor<2x8xui32>, tensor<2x8xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<2x1xui32>, tensor<4xui32>, tensor<4xui32> loc(#loc)
  } loc(#loc76)
  func.func private @_normal(%arg0: tensor<2x2xui32> loc(unknown)) -> tensor<2x8xf32> {
    %0 = call @_normal_real(%arg0) : (tensor<2x2xui32>) -> tensor<2x8xf32> loc(#loc11)
    return %0 : tensor<2x8xf32> loc(#loc)
  } loc(#loc97)
  func.func private @_normal_real(%arg0: tensor<2x2xui32> loc(unknown)) -> tensor<2x8xf32> {
    %cst = stablehlo.constant dense<0x7F800000> : tensor<2x8xf32> loc(#loc13)
    %cst_0 = stablehlo.constant dense<1.000000e+00> : tensor<2x8xf32> loc(#loc13)
    %cst_1 = stablehlo.constant dense<2.83297682> : tensor<2x8xf32> loc(#loc13)
    %cst_2 = stablehlo.constant dense<1.50140941> : tensor<2x8xf32> loc(#loc13)
    %cst_3 = stablehlo.constant dense<1.00167406> : tensor<2x8xf32> loc(#loc13)
    %cst_4 = stablehlo.constant dense<0.246640727> : tensor<2x8xf32> loc(#loc13)
    %cst_5 = stablehlo.constant dense<0.00943887047> : tensor<2x8xf32> loc(#loc13)
    %cst_6 = stablehlo.constant dense<-0.00417768164> : tensor<2x8xf32> loc(#loc13)
    %cst_7 = stablehlo.constant dense<-0.0076224613> : tensor<2x8xf32> loc(#loc13)
    %cst_8 = stablehlo.constant dense<-0.00125372503> : tensor<2x8xf32> loc(#loc13)
    %cst_9 = stablehlo.constant dense<0.00573950773> : tensor<2x8xf32> loc(#loc13)
    %cst_10 = stablehlo.constant dense<2.1858087E-4> : tensor<2x8xf32> loc(#loc13)
    %cst_11 = stablehlo.constant dense<-0.00367342844> : tensor<2x8xf32> loc(#loc13)
    %cst_12 = stablehlo.constant dense<-4.39150654E-6> : tensor<2x8xf32> loc(#loc13)
    %cst_13 = stablehlo.constant dense<0.00134934322> : tensor<2x8xf32> loc(#loc13)
    %cst_14 = stablehlo.constant dense<-3.5233877E-6> : tensor<2x8xf32> loc(#loc13)
    %cst_15 = stablehlo.constant dense<1.00950558E-4> : tensor<2x8xf32> loc(#loc13)
    %cst_16 = stablehlo.constant dense<3.43273939E-7> : tensor<2x8xf32> loc(#loc13)
    %cst_17 = stablehlo.constant dense<-2.00214257E-4> : tensor<2x8xf32> loc(#loc13)
    %cst_18 = stablehlo.constant dense<2.81022636E-8> : tensor<2x8xf32> loc(#loc13)
    %cst_19 = stablehlo.constant dense<3.000000e+00> : tensor<2x8xf32> loc(#loc13)
    %cst_20 = stablehlo.constant dense<2.500000e+00> : tensor<2x8xf32> loc(#loc13)
    %cst_21 = stablehlo.constant dense<5.000000e+00> : tensor<2x8xf32> loc(#loc13)
    %cst_22 = stablehlo.constant dense<1.41421354> : tensor<f32> loc(#loc)
    %cst_23 = stablehlo.constant dense<-0.99999994> : tensor<f32> loc(#loc)
    %cst_24 = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %0 = call @_uniform(%arg0, %cst_23, %cst_24) : (tensor<2x2xui32>, tensor<f32>, tensor<f32>) -> tensor<2x8xf32> loc(#loc14)
    %1 = stablehlo.negate %0 : tensor<2x8xf32> loc(#loc13)
    %2 = stablehlo.multiply %0, %1 : tensor<2x8xf32> loc(#loc13)
    %3 = stablehlo.log_plus_one %2 : tensor<2x8xf32> loc(#loc13)
    %4 = stablehlo.negate %3 : tensor<2x8xf32> loc(#loc13)
    %5 = stablehlo.compare LT, %4, %cst_21 : (tensor<2x8xf32>, tensor<2x8xf32>) -> tensor<2x8xi1> loc(#loc13)
    %6 = stablehlo.subtract %4, %cst_20 : tensor<2x8xf32> loc(#loc13)
    %7 = stablehlo.sqrt %4 : tensor<2x8xf32> loc(#loc13)
    %8 = stablehlo.subtract %7, %cst_19 : tensor<2x8xf32> loc(#loc13)
    %9 = stablehlo.select %5, %6, %8 : tensor<2x8xi1>, tensor<2x8xf32> loc(#loc13)
    %10 = stablehlo.select %5, %cst_18, %cst_17 : tensor<2x8xi1>, tensor<2x8xf32> loc(#loc13)
    %11 = stablehlo.select %5, %cst_16, %cst_15 : tensor<2x8xi1>, tensor<2x8xf32> loc(#loc13)
    %12 = stablehlo.multiply %10, %9 : tensor<2x8xf32> loc(#loc13)
    %13 = stablehlo.add %11, %12 : tensor<2x8xf32> loc(#loc13)
    %14 = stablehlo.select %5, %cst_14, %cst_13 : tensor<2x8xi1>, tensor<2x8xf32> loc(#loc13)
    %15 = stablehlo.multiply %13, %9 : tensor<2x8xf32> loc(#loc13)
    %16 = stablehlo.add %14, %15 : tensor<2x8xf32> loc(#loc13)
    %17 = stablehlo.select %5, %cst_12, %cst_11 : tensor<2x8xi1>, tensor<2x8xf32> loc(#loc13)
    %18 = stablehlo.multiply %16, %9 : tensor<2x8xf32> loc(#loc13)
    %19 = stablehlo.add %17, %18 : tensor<2x8xf32> loc(#loc13)
    %20 = stablehlo.select %5, %cst_10, %cst_9 : tensor<2x8xi1>, tensor<2x8xf32> loc(#loc13)
    %21 = stablehlo.multiply %19, %9 : tensor<2x8xf32> loc(#loc13)
    %22 = stablehlo.add %20, %21 : tensor<2x8xf32> loc(#loc13)
    %23 = stablehlo.select %5, %cst_8, %cst_7 : tensor<2x8xi1>, tensor<2x8xf32> loc(#loc13)
    %24 = stablehlo.multiply %22, %9 : tensor<2x8xf32> loc(#loc13)
    %25 = stablehlo.add %23, %24 : tensor<2x8xf32> loc(#loc13)
    %26 = stablehlo.select %5, %cst_6, %cst_5 : tensor<2x8xi1>, tensor<2x8xf32> loc(#loc13)
    %27 = stablehlo.multiply %25, %9 : tensor<2x8xf32> loc(#loc13)
    %28 = stablehlo.add %26, %27 : tensor<2x8xf32> loc(#loc13)
    %29 = stablehlo.select %5, %cst_4, %cst_3 : tensor<2x8xi1>, tensor<2x8xf32> loc(#loc13)
    %30 = stablehlo.multiply %28, %9 : tensor<2x8xf32> loc(#loc13)
    %31 = stablehlo.add %29, %30 : tensor<2x8xf32> loc(#loc13)
    %32 = stablehlo.select %5, %cst_2, %cst_1 : tensor<2x8xi1>, tensor<2x8xf32> loc(#loc13)
    %33 = stablehlo.multiply %31, %9 : tensor<2x8xf32> loc(#loc13)
    %34 = stablehlo.add %32, %33 : tensor<2x8xf32> loc(#loc13)
    %35 = stablehlo.multiply %34, %0 : tensor<2x8xf32> loc(#loc13)
    %36 = stablehlo.abs %0 : tensor<2x8xf32> loc(#loc13)
    %37 = stablehlo.compare EQ, %36, %cst_0 : (tensor<2x8xf32>, tensor<2x8xf32>) -> tensor<2x8xi1> loc(#loc13)
    %38 = stablehlo.multiply %0, %cst : tensor<2x8xf32> loc(#loc13)
    %39 = stablehlo.select %37, %38, %35 : tensor<2x8xi1>, tensor<2x8xf32> loc(#loc13)
    %40 = stablehlo.broadcast_in_dim %cst_22, dims = [] : (tensor<f32>) -> tensor<2x8xf32> loc(#loc15)
    %41 = stablehlo.multiply %40, %39 : tensor<2x8xf32> loc(#loc15)
    return %41 : tensor<2x8xf32> loc(#loc)
  } loc(#loc32)
  func.func private @_uniform(%arg0: tensor<2x2xui32> loc(unknown), %arg1: tensor<f32> loc(unknown), %arg2: tensor<f32> loc(unknown)) -> tensor<2x8xf32> {
    %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32> loc(#loc)
    %c = stablehlo.constant dense<1065353216> : tensor<ui32> loc(#loc)
    %c_0 = stablehlo.constant dense<9> : tensor<ui32> loc(#loc)
    %c_1 = stablehlo.constant dense<32> : tensor<ui64> loc(#loc59)
    %c_2 = stablehlo.constant dense<1> : tensor<ui64> loc(#loc59)
    %0 = stablehlo.broadcast_in_dim %arg1, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc16)
    %1 = stablehlo.broadcast_in_dim %arg2, dims = [] : (tensor<f32>) -> tensor<1xf32> loc(#loc16)
    %2 = stablehlo.slice %arg0 [0:2, 0:1] : (tensor<2x2xui32>) -> tensor<2x1xui32> loc(#loc60)
    %3 = stablehlo.reshape %2 : (tensor<2x1xui32>) -> tensor<2xui32> loc(#loc60)
    %4 = stablehlo.slice %arg0 [0:2, 1:2] : (tensor<2x2xui32>) -> tensor<2x1xui32> loc(#loc60)
    %5 = stablehlo.reshape %4 : (tensor<2x1xui32>) -> tensor<2xui32> loc(#loc60)
    %6 = stablehlo.iota dim = 0 : tensor<8xui64> loc(#loc59)
    %7 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<ui64>) -> tensor<8xui64> loc(#loc59)
    %8 = stablehlo.multiply %7, %6 : tensor<8xui64> loc(#loc59)
    %9 = stablehlo.broadcast_in_dim %c_1, dims = [] : (tensor<ui64>) -> tensor<8xui64> loc(#loc59)
    %10 = stablehlo.shift_right_logical %8, %9 : tensor<8xui64> loc(#loc59)
    %11 = stablehlo.convert %8 : (tensor<8xui64>) -> tensor<8xui32> loc(#loc59)
    %12 = stablehlo.convert %10 : (tensor<8xui64>) -> tensor<8xui32> loc(#loc59)
    %13 = stablehlo.broadcast_in_dim %12, dims = [1] : (tensor<8xui32>) -> tensor<1x8xui32> loc(#loc61)
    %14 = stablehlo.broadcast_in_dim %11, dims = [1] : (tensor<8xui32>) -> tensor<1x8xui32> loc(#loc61)
    %15 = stablehlo.broadcast_in_dim %3, dims = [0] : (tensor<2xui32>) -> tensor<2x1xui32> loc(#loc61)
    %16 = stablehlo.broadcast_in_dim %5, dims = [0] : (tensor<2xui32>) -> tensor<2x1xui32> loc(#loc61)
    %17:2 = call @threefry2x32_0(%15, %16, %13, %14) : (tensor<2x1xui32>, tensor<2x1xui32>, tensor<1x8xui32>, tensor<1x8xui32>) -> (tensor<2x8xui32>, tensor<2x8xui32>) loc(#loc48)
    %18 = stablehlo.xor %17#0, %17#1 : tensor<2x8xui32> loc(#loc62)
    %19 = stablehlo.broadcast_in_dim %c_0, dims = [] : (tensor<ui32>) -> tensor<2x8xui32> loc(#loc17)
    %20 = stablehlo.shift_right_logical %18, %19 : tensor<2x8xui32> loc(#loc17)
    %21 = stablehlo.broadcast_in_dim %c, dims = [] : (tensor<ui32>) -> tensor<2x8xui32> loc(#loc18)
    %22 = stablehlo.or %20, %21 : tensor<2x8xui32> loc(#loc18)
    %23 = stablehlo.bitcast_convert %22 : (tensor<2x8xui32>) -> tensor<2x8xf32> loc(#loc19)
    %24 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<2x8xf32> loc(#loc20)
    %25 = stablehlo.subtract %23, %24 : tensor<2x8xf32> loc(#loc20)
    %26 = stablehlo.subtract %1, %0 : tensor<1xf32> loc(#loc20)
    %27 = stablehlo.broadcast_in_dim %26, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc16)
    %28 = stablehlo.broadcast_in_dim %27, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<2x8xf32> loc(#loc15)
    %29 = stablehlo.multiply %25, %28 : tensor<2x8xf32> loc(#loc15)
    %30 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc16)
    %31 = stablehlo.broadcast_in_dim %30, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<2x8xf32> loc(#loc21)
    %32 = stablehlo.add %29, %31 : tensor<2x8xf32> loc(#loc21)
    %33 = stablehlo.broadcast_in_dim %0, dims = [1] : (tensor<1xf32>) -> tensor<1x1xf32> loc(#loc16)
    %34 = stablehlo.broadcast_in_dim %33, dims = [0, 1] : (tensor<1x1xf32>) -> tensor<2x8xf32> loc(#loc22)
    %35 = stablehlo.maximum %34, %32 : tensor<2x8xf32> loc(#loc22)
    return %35 : tensor<2x8xf32> loc(#loc)
  } loc(#loc32)
} loc(#loc)
#loc3 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5093/prefile.py":48:13 to 50:3)
#loc4 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5093/prefile.py":114:2 to :8)
#loc6 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5093/prefile.py":34:2 to :8)
#loc7 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5093/prefile.py":35:14 to :34)
#loc8 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5093/prefile.py":35:2 to :11)
#loc9 = loc("/home/marielle/.claude/jobs/a73ad07c/tmp/probe5093/prefile.py":38:4 to :50)
#loc10 = loc("threefry2x32")
#loc11 = loc("jit(_normal_real)")
#loc12 = loc("jit")
#loc13 = loc("erf_inv")
#loc14 = loc("jit(_uniform)")
#loc15 = loc("mul")
#loc16 = loc("broadcast_in_dim")
#loc17 = loc("shift_right_logical")
#loc18 = loc("or")
#loc19 = loc("bitcast_convert_type")
#loc20 = loc("sub")
#loc21 = loc("add")
#loc22 = loc("max")
#loc24 = loc("main"(#loc3))
#loc25 = loc("<module>"(#loc4))
#loc27 = loc("minimal"(#loc6))
#loc28 = loc("minimal"(#loc7))
#loc29 = loc("minimal"(#loc8))
#loc30 = loc("minimal"(#loc9))
#loc31 = loc("threefry2x32:"(#loc10))
#loc32 = loc("jit:"(#loc12))
#loc33 = loc(callsite(#loc24 at #loc25))
#loc34 = loc("vmap()/jit"(#loc26))
#loc35 = loc("iota_2x32_shape"(#loc26))
#loc36 = loc("unstack"(#loc26))
#loc37 = loc("broadcast_in_dim"(#loc26))
#loc38 = loc("stack"(#loc26))
#loc39 = loc("xor"(#loc26))
#loc40 = loc("add"(#loc26))
#loc41 = loc(callsite(#loc26 at #loc26))
#loc42 = loc("while/body/eval_jaxpr"(#loc26))
#loc43 = loc("shift_left"(#loc26))
#loc44 = loc("sub"(#loc26))
#loc45 = loc("shift_right_logical"(#loc26))
#loc46 = loc("or"(#loc26))
#loc47 = loc("convert_element_type"(#loc26))
#loc49 = loc("xor"(#loc23))
#loc50 = loc("add"(#loc23))
#loc51 = loc(callsite(#loc23 at #loc23))
#loc52 = loc("while/body/eval_jaxpr"(#loc23))
#loc53 = loc("unstack"(#loc23))
#loc54 = loc("shift_left"(#loc23))
#loc55 = loc("sub"(#loc23))
#loc56 = loc("shift_right_logical"(#loc23))
#loc57 = loc("or"(#loc23))
#loc58 = loc("convert_element_type"(#loc23))
#loc59 = loc("vmap()/iota_2x32_shape"(#loc23))
#loc60 = loc("vmap()/unstack"(#loc23))
#loc61 = loc("vmap()/broadcast_in_dim"(#loc23))
#loc62 = loc("vmap()/xor"(#loc23))
#loc63 = loc(callsite(#loc23 at #loc33))
#loc64 = loc(callsite(#loc26 at #loc33))
#loc65 = loc(callsite(#loc27 at #loc33))
#loc66 = loc(callsite(#loc28 at #loc33))
#loc67 = loc(callsite(#loc29 at #loc33))
#loc68 = loc(callsite(#loc30 at #loc33))
#loc69 = loc("jit(minimal)"(#loc33))
#loc70 = loc("jit:"(#loc34))
#loc71 = loc("while"(#loc41))
#loc72 = loc(callsite(#loc41 at #loc26))
#loc73 = loc("eval_jaxpr:"(#loc42))
#loc74 = loc("while"(#loc51))
#loc75 = loc(callsite(#loc51 at #loc23))
#loc76 = loc("eval_jaxpr:"(#loc52))
#loc77 = loc(callsite(#loc23 at #loc63))
#loc78 = loc(callsite(#loc26 at #loc64))
#loc79 = loc("jit(minimal)/vmap()/unstack"(#loc65))
#loc80 = loc(callsite(#loc26 at #loc66))
#loc81 = loc("jit(minimal)/vmap()/unstack"(#loc67))
#loc82 = loc("jit(minimal)/vmap(jit(_normal))"(#loc68))
#loc83 = loc("while/cond/lt"(#loc72))
#loc84 = loc("while/body/eval_jaxpr"(#loc72))
#loc85 = loc("while/body/add"(#loc72))
#loc86 = loc("while/cond/lt"(#loc75))
#loc87 = loc("while/body/eval_jaxpr"(#loc75))
#loc88 = loc("while/body/add"(#loc75))
#loc89 = loc("jit(minimal)/vmap()/jit"(#loc68))
#loc90 = loc("jit(minimal)/vmap()/vmap()/iota_2x32_shape"(#loc77))
#loc91 = loc("jit(minimal)/vmap()/vmap(jit(_threefry_split))"(#loc78))
#loc92 = loc("jit(minimal)/vmap()/vmap(jit(_threefry_split))"(#loc80))
#loc93 = loc("jit(minimal)/vmap()/vmap()/unstack"(#loc77))
#loc94 = loc("jit(minimal)/vmap()/vmap()/broadcast_in_dim"(#loc77))
#loc95 = loc("jit(minimal)/vmap()/vmap()"(#loc77))
#loc96 = loc("jit(minimal)/vmap()/vmap()/xor"(#loc77))
#loc97 = loc("jit:"(#loc89))

What component(s) does this issue relate to?

Compiler, MLIR

Version information

iree-base-compiler 3.11.0   (3.11.0rc20260316 @ e4a3b0405d7d23554da26403658d0e8c3c5ecf25)
iree-base-runtime  3.11.0
LLVM               23.0.0git
jax / jaxlib       0.11.1
host               Linux x86-64 (WSL2)

Additional context

Triage already done — this is not CPU-codegen-specific:

variation	result
--iree-llvmcpu-target-cpu=host	wrong, as above
default flags (no target-cpu)	wrong, byte-identical to the above
--iree-hal-target-backends=vmvx	wrong, byte-identical to the above
--iree-opt-const-eval=false	wrong, byte-identical

The reference vmvx backend producing the same wrong bits as llvm-cpu suggests the defect is above backend codegen, somewhere in the shared pipeline. The output is fully deterministic across runs and across flag sets.

A second, possibly related defect. The identical program without jax.vmap (one bare call, uint32[2] input rather than 2x2) does not return wrong values — it aborts the runtime with what looks like a corrupt buffer length:

IndexError: Error invoking function:
  c/runtime/src/iree/hal/command_buffer_validation.c:143: OUT_OF_RANGE;
  at least one command attempted to access an address outside of the valid
  bound buffer range (length=549755813960, end(inc)=549755813959,
  binding offset=0, binding length=192, binding end(inc)=191);
  binding table slot 0; while invoking native function
  hal.device.queue.execute.indirect

length=549755813960 is 0x8000000C8 against a 192-byte binding. Happy to split that into its own issue if you'd prefer — I kept it here because it reproduces on the same tiny program with only the vmap removed, so the two may share a cause.

Why this is worth catching: the divergence is carried by integer values, so a float-tolerance parity check between the JAX and IREE outputs reports max_abs_diff == 0.0 on the integer output and passes. In our case it surfaced as a model's permutation coming back as a different valid permutation — no crash, no NaN, nothing out of range, just quietly different sampling. We only found it because an unrelated extra read of the key made the program exact again.

I have a tracked reproducer script carrying four negative controls (it fails loudly if any control stops holding) and am happy to run further variations, bisect flags, or test a patch on request.

---

AI tool use disclosure

Adding this to comply with the project's
AI tool use policy,
which I should have included when opening this issue rather than after — apologies.

Assisted-by: Claude Opus 5 (Claude Code)

The investigation, the minimization, and the text of this report were substantially
AI-assisted. Everything quoted above is machine-measured rather than asserted: the numbers
come from the commands shown, and the whole thing is re-derivable from the attached MLIR
with iree-compile + iree-run-module alone. The Expected values were cross-checked
against XLA running the same module.

Happy to trim, restructure, or re-write any of this in my own words if you'd prefer a
shorter report — just say so.

[github-actions]

Hello @maraxen 👋

Thank you for submitting an issue to IREE! It looks like it's your first one.

The volume of issues is huge, and triage times may vary. If you need quicker assistance on the issue, consider posting it on IREE Discord and providing short context there.
Feel free to identify the most frequent contributors to the corresponding part of the codebase via Git history/CODEOWNERS and (gently!) tag them below.

---

Action required: acknowledge IREE project policies

IREE is a Linux Foundation project. All participants are expected to follow the LF Projects Code of Conduct.

Please note that all contributions to IREE, including GitHub issues, must follow our IREE AI Tool Use Policy. This means:

• Issue reporters must fully understand the primary issue and the impact of the problem.
  For bug reports, specifics of the use-case, workflow and/or crash scenario usually suffice.
  For improvement suggestions, this usually implies deeper understanding of suboptimal code, IR design choices, assumed performance tradeoffs, etc.
• Contributors must write the summary of the issue themselves, so that the exact impact and scenario is clear to other participants.
• Issue authors are expected to review and understand all LLM/AI-assisted segments of the report, if any. E.g.: CLI commands and IR reproducers, model/framework details, lowering pipeline details, assumptions about problematic behaviors of IREE components.
• Substantial use of LLM/generative AI tools must be disclaimed in the header or footer of the issue description.

Before the issue can be triaged, we kindly ask you to reply to this message and confirm that you understand and accept the proposed policy. As required, please adjust your initial issue description.

---

If you have any questions, feel free to leave a comment here, or ask away on IREE Discord.

Thank you,
The IREE Community

[maraxen]

Splitting out the runtime abort I mentioned under Additional context here: it
is now #24929.

Having characterised it, I no longer think it is the same defect as this issue:

	this issue (#24927)	#24929
symptom	silently wrong values	hard abort at invocation
trigger	a nested split feeding a draw, with a second draw live	un-batched threefry on a runtime-input key
jax_threefry_partitionable=False	still reproduces	fixes it
constant key	still reproduces	exact

The partitionable row is what convinced me they are independent. This issue
stands on its own and nothing about it has changed.

Assisted-by: Claude Opus 5 (Claude Code)