Repository: jax-ml/jax (python)

Issue:
expi(x) returns incorrect value when jit disabled
### Description

Specific case when incorrect result is returned:

- Input is a one-element array, valued in range of `-1 <= x < 0.`
- `jax_disable_jit` is set to `True`

Simple reproducible code:
```
import jax.numpy as jnp
from jax.scipy.special import expi

jax.config.update("jax_disable_jit", False)
expi(jnp.array([-0.5])) # correct result

jax.config.update("jax_disable_jit", True)
expi(jnp.array([-0.5])) # incorrect result
```

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

```
jax:    0.5.1
jaxlib: 0.5.1
numpy:  2.2.2
python: 3.12.2 | packaged by conda-forge | (main, Feb 16 2024, 20:54:21) [Clang 16.0.6 ]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node='MacBook-Pro', release='23.6.0', version='Darwin Kernel Version 23.6.0: Mon Jul 29 21:14:04 PDT 2024; root:xnu-10063.141.2~1/RELEASE_ARM64_T8122', machine='arm64')
```

Pull request description:
jax.disable_jit: ensure while_loop behaves similarly to non-disable_jit version
Fixes #27019

Patch:
diff --git a/jax/_src/lax/control_flow/loops.py b/jax/_src/lax/control_flow/loops.py
index c7dee3e71..3084fa722 100644
--- a/jax/_src/lax/control_flow/loops.py
+++ b/jax/_src/lax/control_flow/loops.py
@@ -1364,9 +1364,9 @@ def while_loop(cond_fun: Callable[[T], BooleanNumeric],
     raise TypeError("lax.while_loop: body_fun and cond_fun arguments should be callable.")
   if config.disable_jit.value:
     try:
-      val = init_val
+      val = tree_map(lax.asarray, init_val)
       while cond_fun(val):
-        val = body_fun(val)
+        val = tree_map(lax.asarray, body_fun(val))
       return val
     except core.ConcretizationTypeError:
       # Can't run this while_loop in Python (e.g. because there's a vmap
diff --git a/jax/experimental/mosaic/gpu/core.py b/jax/experimental/mosaic/gpu/core.py
index 66e19bb5f..c92074dc2 100644
--- a/jax/experimental/mosaic/gpu/core.py
+++ b/jax/experimental/mosaic/gpu/core.py
@@ -260,7 +260,7 @@ def _construct_smem_reftree(
             dynamic_smem, c(dynamic_smem_offset, index), [],
         )
         if layout is None:
-          layout = tcgen05._infer_tmem_layout(shape, collective)
+          layout = tcgen05._infer_tmem_layout(shape)
         num_cols = layout.cols_in_shape(shape)
         delayed_warp_init.append(
             functools.partial(
diff --git a/jax/experimental/mosaic/gpu/examples/matmul_blackwell.py b/jax/experimental/mosaic/gpu/examples/matmul_blackwell.py
index 6af394d00..d15cecbdc 100644
--- a/jax/experimental/mosaic/gpu/examples/matmul_blackwell.py
+++ b/jax/experimental/mosaic/gpu/examples/matmul_blackwell.py
@@ -230,8 +230,8 @@ def main(unused_argv):
       tile_n *= 2
     if m < tile_m or n < tile_n:
       continue
-    if tile_n > 512:
-      continue
+    if kwargs["collective"] and tile_n >= 512:
+      continue  # TODO(apaszke): Support 512
     if (m // tile_m) % kwargs["grid_tile_m"]:
       continue
     try:
diff --git a/jax/experimental/mosaic/gpu/tcgen05.py b/jax/experimental/mosaic/gpu/tcgen05.py
index e5a2d3aa5..7a349f50c 100644
--- a/jax/experimental/mosaic/gpu/tcgen05.py
+++ b/jax/experimental/mosaic/gpu/tcgen05.py
@@ -83,7 +83,6 @@ def mma(
     accumulate: ir.Value | bool = True,
     collective: bool = False,
 ):
-  i32 = ir.IntegerType.get_signless(32)
   i64 = ir.IntegerType.get_signless(64)
   if isinstance(accumulate, bool):
     accumulate = arith.constant(ir.IntegerType.get_signless(1), accumulate)
@@ -113,10 +112,6 @@ def mma(
     raise ValueError(
         f"Accumulator shape mismatch: expected {(m, n * num_cta)}, got {d.shape}"
     )
-  if d.layout != (expected_layout := _infer_tmem_layout(d.shape, collective)):
-    raise ValueError(
-        f"Accumulator layout mismatch: expected {expected_layout}, got {d.layout}"
-    )
   f32 = ir.F32Type.get()
   if element_type == f32 or element_type == ir.BF16Type.get():
     if d.dtype != f32:
@@ -141,7 +136,11 @@ def mma(
     raise ValueError(f"N must be a multiple of 8, got: {n}")
   elif n > 256 and n != 512:
     raise ValueError("Only N below 256 or N=512 are supported")
-  n_group_elems = min(n, 256 // num_cta)
+  if num_cta == 2 and n > 256:
+    raise NotImplementedError(
+        "N is too big for collective MMA. Only up to 256 is supported."
+    )
+  n_group_elems = min(n, 256)
   if m % m_group_elems:
     raise ValueError(f"M must be a multiple of {m_group_elems}, got: {m}")
   if k % k_group_elems:
@@ -180,7 +179,6 @@ def mma(
 
   # Step 4. Issue the instructions.
   true = arith.constant(ir.IntegerType.get_signless(1), 1)
-  n_collective_group_elems = n_group_elems * num_cta
   for mi, ni, ki in np.ndindex(m_groups, n_groups, k_groups):
     a_offset = mi * a_m_group_stride + ki * a_k_group_stride
     a_mk = arith.addi(a_desc_base, utils.c(mma_utils.encode_addr(a_offset), i64))
@@ -190,9 +188,9 @@ def mma(
       raise NotImplementedError("D needs to be sliced")
     acc = accumulate if ki == 0 else true
     _do_mma(
-        arith.addi(
-            d.address, arith.constant(i32, ni * n_collective_group_elems)
-        ),
+        d.slice(
+            slice(None), utils.ds(ni * n_group_elems, n_group_elems)
+        ).address,
         a_mk,
         b_nk,
         d_type=ir.F32Type.get(),
@@ -379,15 +377,8 @@ class TMEMLayout:
     +------------------+------------------+
     | [0:64, 64:128]   | [64:128, 64:128] |
     +------------------+------------------+
-
-  The above is further complicated by column_tile_stride, which is used to
-  swizzle the ordering of column tiles. That is, if column_tile_stride is 2,
-  we will first lay out all tiles that have the column index 0, 2, 4, and so on
-  until we run out of tiles. Only then we lay out the tiles with column index
-  1, 3, etc.
   """
   elements_in_tile: tuple[int, int]
-  column_tile_stride: int = 1
 
   def __post_init__(self):
     row_tiling = self.elements_in_tile[0]
@@ -414,7 +405,7 @@ class TMEMLayout:
     return num_tiles // tiles_in_row * cols_in_tile
 
 
-def _infer_tmem_layout(shape: tuple[int, int], collective: bool) -> TMEMLayout:
+def _infer_tmem_layout(shape: tuple[int, int]) -> TMEMLayout:
   if shape[0] > TMEM_ROWS:
     raise ValueError(
         "Can only infer TMEM layout for shapes with at most 128 rows, got:"
@@ -430,15 +421,7 @@ def _infer_tmem_layout(shape: tuple[int, int], collective: bool) -> TMEMLayout:
         "Can only infer TMEM layout for shapes with row count that's a power of"
         f" 2, got: {shape[0]}"
     )
-  if shape[1] % 8:
-    raise ValueError(
-        "Can only infer TMEM layout for shapes with column count that's a"
-        f" multiple of 8, got: {shape[1]}"
-    )
-  if collective and shape[1] == 512:
-    return TMEMLayout(elements_in_tile=(shape[0], 128), column_tile_stride=2)
-  else:
-    return TMEMLayout(elements_in_tile=(shape[0], 8))
+  return TMEMLayout(elements_in_tile=(shape[0], 1))
 
 
 @dataclasses.dataclass(frozen=True)
@@ -449,14 +432,7 @@ class TMEMRef:
   layout: TMEMLayout
 
   @classmethod
-  def from_alloc(
-      cls,
-      tmem_addr_ref: ir.Value,
-      shape: tuple[int, int],
-      dtype,
-      collective: bool | None = None,
-      layout: TMEMLayout | None = None,
-  ):
+  def from_alloc(cls, tmem_addr_ref: ir.Value, shape: tuple[int, int], dtype, layout: TMEMLayout | None = None):
     i32 = ir.IntegerType.get_signless(32)
     if not ir.MemRefType.isinstance(tmem_addr_ref.type):
       raise ValueError(f"tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}")
@@ -473,11 +449,7 @@ class TMEMRef:
     if shape[0] < 32:
       raise ValueError(f"TMEM refs must have at least 32 rows, got: {shape[0]}")
     if layout is None:
-      if collective is None:
-        raise ValueError(
-            "collective argument must be provided when TMEM layout is inferred"
-        )
-      layout = _infer_tmem_layout(shape, collective)
+      layout = _infer_tmem_layout(shape)
     else:
       layout.check_shape(shape)
     # TODO: Do we have to do this??
@@ -489,17 +461,12 @@ class TMEMRef:
     base_idx, slice_shape, is_squeezed = utils.parse_indices(idxs, self.shape)
     if any(is_squeezed):
       raise ValueError("TMEM can only be sliced, not indexed")
-    if self.layout != TMEMLayout(elements_in_tile=(TMEM_ROWS, 8)):
+    if self.layout.elements_in_tile[0] != TMEM_ROWS:
       raise NotImplementedError(
-          "Slicing only implemented for refs with standard layout, got:"
-          f" {self.layout}"
+          f"Slicing only implemented for refs with tiling of {TMEM_ROWS} rows"
       )
     if base_idx[0] != 0 or slice_shape[0] != TMEM_ROWS:
       raise NotImplementedError("TMEM cannot be sliced along rows")
-    if slice_shape[1] % 8:
-      raise NotImplementedError(
-          "TMEM column slice length must be a multiple of 8"
-      )
     col_idx = base_idx[1]
     if not isinstance(col_idx, ir.Value):
       col_idx = arith.constant(ir.IntegerType.get_signless(32), col_idx)
@@ -517,75 +484,48 @@ class TMEMRef:
       raise ValueError("TMEM loads only support slicing")
     if any(idx != 0 for idx in base_idxs) or tuple(slice_shape) != self.shape:
       raise NotImplementedError("Slicing of TMEM not impelmented yet")
+    if self.layout.elements_in_tile[0] != TMEM_ROWS:
+      raise NotImplementedError(
+          f"Loads only implemented for refs with tiling of {TMEM_ROWS} rows"
+      )
     if self.shape[1] % 8:
       raise NotImplementedError
     if self.dtype != ir.F32Type.get():
       raise NotImplementedError(self.dtype)
     layout = _m128_256bit_32bit_layout(self.shape)
     regs_shape = layout.registers_shape(self.shape)
-    if self.layout == TMEMLayout(elements_in_tile=(TMEM_ROWS, 8)):
-      # load_32xcols returns a 4xN array, but the FA tiling we use here tiles
-      # columns before rows, and so it is Nx4 (after ignoring all 1 dims).
-      registers = _load_32xcols(
-          self.address, self.shape[1], self.dtype
-      ).T.reshape(regs_shape)
-    elif self.layout == TMEMLayout(elements_in_tile=(TMEM_ROWS, 128), column_tile_stride=2):
-      if self.shape[1] % 128 != 0:
-        raise ValueError(
-            f"TMEM layout {self.layout} is not compatible with shape {self.shape}"
-        )
-      num_column_tiles = self.shape[1] // 128
-      column_tile_stride = self.layout.column_tile_stride
-      num_strided_col_groups = utils.ceil_div(num_column_tiles, column_tile_stride)
-      tiles = []
-      for col_tile_base in range(num_strided_col_groups):
-        for col_tile in range(col_tile_base, num_column_tiles, column_tile_stride):
-          tiles.append(
-              _load_32xcols(
-                  arith.addi(self.address, arith.constant(i32, col_tile * 128)),
-                  cols=128,
-                  dtype=self.dtype,
-              )
-          )
-      registers = np.concatenate(tiles, axis=1).T.reshape(regs_shape)
+    num = self.shape[1] // 8
+    # TODO(apaszke): Make the tiling configurable through the args too.
+    if num <= 32:
+      num_tiling = num
+    elif num == 64:
+      num_tiling = 32
     else:
-      raise NotImplementedError(
-          f"Loads only implemented for refs with standard layout, got: {self.layout}"
-      )
+      raise NotImplementedError(num)
+    registers = np.empty(regs_shape, dtype=object)
+    # We load 16 lanes at a time, but need 32 in total.
+    for row_group in range(2):
+      addr_row = arith.addi(self.address, arith.constant(i32, (row_group * 16) << 16))
+      regs = []
+      cols_per_num_tile = 8  # This depends on the 16x256b below.
+      for num_group in range(num // num_tiling):
+        addr_row_col = arith.addi(
+            addr_row,
+            arith.constant(i32, num_tiling * num_group * cols_per_num_tile),
+        )
+        regs += tmem_load(addr_row_col, "16x256b", num_tiling)
+      regs = [llvm.bitcast(self.dtype, r) for r in regs]
+      vector_regs = []
+      undef = llvm.mlir_undef(ir.VectorType.get((2,), self.dtype))
+      for r_low, r_high in zip(regs[::2], regs[1::2]):
+        high_undef = llvm.insertelement(undef, r_low, utils.c(0, i32))
+        vreg = llvm.insertelement(high_undef, r_high, utils.c(1, i32))
+        vector_regs.append(vreg)
+      # Dimension 4 is the one where we split 32 rows into tiles of 8.
+      regs_slice = (slice(None),) * 4 + (slice(row_group * 2, (row_group + 1) * 2),)
+      registers[regs_slice] = np.asarray(vector_regs, dtype=object).reshape(registers[regs_slice].shape)
     return fa.FragmentedArray(_registers=registers, _layout=layout, _is_signed=None)
 
-def _load_32xcols(base_addr, cols, dtype):
-  # See https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-matrix-fragments-shape-16256b
-  i32 = ir.IntegerType.get_signless(32)
-  assert cols % 8 == 0
-  cols_per_num_tile = 8
-  load_shape = "16x256b"
-  num = cols // 8
-  if num <= 32:
-    num_tiling = num
-  elif num == 64:
-    num_tiling = 32
-  else:
-    raise NotImplementedError(num)
-  vector_regs = np.ndarray((4, num), dtype=object)
-  # We load 16 lanes at a time, but need 32 in total.
-  for row_group in range(2):
-    addr_row = arith.addi(base_addr, arith.constant(i32, (row_group * 16) << 16))
-    regs = []
-    for num_group in range(num // num_tiling):
-      addr_row_col = arith.addi(
-          addr_row,
-          arith.constant(i32, num_tiling * num_group * cols_per_num_tile),
-      )
-      regs += tmem_load(addr_row_col, load_shape, num_tiling)
-    regs = [llvm.bitcast(dtype, r) for r in regs]
-    undef = llvm.mlir_undef(ir.VectorType.get((2,), dtype))
-    for r_low, r_high, idx in zip(regs[::2], regs[1::2], np.ndindex(num, 2)):
-      high_undef = llvm.insertelement(undef, r_low, utils.c(0, i32))
-      vreg = llvm.insertelement(high_undef, r_high, utils.c(1, i32))
-      vector_regs[idx[1] + 2 * row_group, idx[0]] = vreg
-  return vector_regs
-
 
 def _m128_256bit_32bit_layout(shape: tuple[int, ...]):
   if len(shape) != 2:
diff --git a/jax/experimental/mosaic/gpu/utils.py b/jax/experimental/mosaic/gpu/utils.py
index 080397bbb..f90f7ff08 100644
--- a/jax/experimental/mosaic/gpu/utils.py
+++ b/jax/experimental/mosaic/gpu/utils.py
@@ -1201,7 +1201,3 @@ def bitcast(x: ir.Value, new_type: ir.Type):
     assert x_ty.width == bitwidth(new_type.element_type) * math.prod(new_type.shape)
     return vector.bitcast(new_type, vector.splat(ir.VectorType.get((1,), x_ty), x))
   raise ValueError(f"Can't bitcast {x.type} to {new_type}")
-
-
-def ceil_div(x: int, y: int):
-  return (x + y - 1) // y


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
index 15fc37805..39da45079 100644
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -2502,6 +2502,28 @@ class LaxControlFlowTest(jtu.JaxTestCase):
     x, n = jnp.arange(3), jnp.arange(4)
     jax.vmap(jax.vmap(f, (None, 0)), (0, None))(x, n)  # doesn't crash
 
+  def test_disable_jit_while_loop_with_mutation(self):
+    # https://github.com/jax-ml/jax/issues/27019
+
+    def body_fun(carry):
+      x, y = carry
+      x += 1  # in-place if x is mutable
+      return x, y + x
+
+    def cond_fun(carry):
+      x, _ = carry
+      return x < 10
+
+    def f():
+      val = np.array(1.0)  # mutable value
+      return jax.lax.while_loop(cond_fun, body_fun, (val, val))[1]
+
+    with jax.disable_jit(False):
+      result_jit = f()
+    with jax.disable_jit(True):
+      result_nojit = f()
+    self.assertEqual(result_jit, result_nojit)
+
   @parameterized.named_parameters(
       {"testcase_name": f"_{shape}_{axis=}",
        "shape": shape, "axis": axis}
diff --git a/tests/lax_scipy_special_functions_test.py b/tests/lax_scipy_special_functions_test.py
index 96d48dcd3..8932869f1 100644
--- a/tests/lax_scipy_special_functions_test.py
+++ b/tests/lax_scipy_special_functions_test.py
@@ -278,6 +278,15 @@ class LaxScipySpcialFunctionsTest(jtu.JaxTestCase):
     with jax.checking_leaks():
       lsp_special.expi(jnp.ones(()))
 
+  def testExpiDisableJit(self):
+    # Regression test for https://github.com/jax-ml/jax/issues/27019
+    x = jnp.array([-0.5])
+    with jax.disable_jit(True):
+      result_nojit = lsp_special.expi(x)
+    with jax.disable_jit(False):
+      result_jit = lsp_special.expi(x)
+    self.assertAllClose(result_jit, result_nojit)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
diff --git a/tests/mosaic/gpu_test.py b/tests/mosaic/gpu_test.py
index f6b94b777..cc654eb2b 100644
--- a/tests/mosaic/gpu_test.py
+++ b/tests/mosaic/gpu_test.py
@@ -1026,7 +1026,7 @@ class TCGen05Test(TestCase):
       in_jax_dtype=(jnp.float16,),  # TODO(apaszke): f32
       out_jax_dtype=(jnp.float32,),  # TODO(apaszke): f16 accumulation
       m=(256,),  # TODO(apaszke): 64, 192, 256
-      n=(128, 256, 512),  # TODO(apaszke): 192, other non-power-of-2
+      n=(128, 256),  # TODO(apaszke): 512, 192, other non-power-of-2
       k_steps=(1, 2),
       swizzle=(32, 64, 128,),
   )