[Pallas] Raise NotImplementedError for strided load/store in interpret mode.

PiperOrigin-RevId: 615983065

Author: bythew3i

jax/_src/pallas/primitives.py

@@ -303,6 +303,10 @@ def _load_discharge_rule(in_avals, out_avals, *args_flat, args_tree, **_):
     raise NotImplementedError("Only one indexer supported in discharge rule.")
   idx = indexers[0]
   if all((isinstance(s, Slice) or not s.shape) for s in idx.indices):
+    # TODO(b/329733289): support strided load/store in interpret mode.
+    for s in idx.indices:
+      if isinstance(s, Slice) and s.stride > 1:
+        raise NotImplementedError("Unimplemented stride support.")
     indices = idx.indices
     scalar_dims = [not isinstance(s, Slice) and not s.shape for s in indices]
     slice_starts = [s.start if isinstance(s, Slice) else s for s in indices]
@@ -404,6 +408,10 @@ def _swap_discharge_rule(in_avals, out_avals, *args_flat, args_tree, **_):
     raise NotImplementedError("Only one indexer supported in discharge rule.")
   idx = indexers[0]
   if all((isinstance(s, Slice) or not s.shape) for s in idx.indices):
+    # TODO(b/329733289): support strided load/store in interpret mode.
+    for s in idx.indices:
+      if isinstance(s, Slice) and s.stride > 1:
+        raise NotImplementedError("Unimplemented stride support.")
     indices = idx.indices
     scalar_dims = [
         i

jax/_src/state/discharge.py

@@ -194,6 +194,10 @@ def _maybe_convert_to_dynamic_slice(
   if not all(isinstance(i, indexing.Slice) or not np.shape(i)
              for i in indexer.indices):
     return None
+  # TODO(b/329733289): support strided load/store in interpret mode.
+  for i in indexer.indices:
+    if isinstance(i, indexing.Slice) and i.stride > 1:
+      raise NotImplementedError("Unimplemented stride support.")
   _convert_i32 = lambda x: lax.convert_element_type(x, np.dtype("int32"))
   starts = tuple(
       _convert_i32(i.start) if isinstance(i, indexing.Slice)