Fix Ref indexing discrepancies with array indexing

This commit addresses issues where Ref indexing behaved differently
from JAX array indexing:

1. OOB slice clamping: Allow start == dim when size == 0 (empty slice)
   instead of raising an error

2. None indexing support: Add RefNewAxis transform to handle np.newaxis
   (None) in indices, enabling patterns like x[None] and x[..., None]

3. Improved error message for negative slice steps to clarify this is
   a known limitation of Ref indexing

Fixes parts of issue #33322.

Author: yashwantbezawada

jax/_src/state/indexing.py

@@ -75,7 +75,10 @@ def tree_unflatten(cls, aux_data, children) -> Slice:
   def from_slice(cls, slc: slice, size: int) -> Slice:
     start, step, size = core.canonicalize_slice(slc, size)
     if step < 1:
-      raise ValueError(f"slice must have a step >= 1 (found: {step})")
+      raise ValueError(
+          f"Ref indexing does not support negative slice steps (found: {step}). "
+          "Consider using positive steps and reversing the result if needed."
+      )
     return cls(start, size, step)
 
 
@@ -172,12 +175,14 @@ def __post_init__(self):
       if isinstance(idx, Slice):
         start = idx.start
         if value := _maybe_concretize(start):
-          if value >= s:
+          size_val = _maybe_concretize(idx.size)
+          # Allow start == s when size == 0 (empty slice)
+          if value >= s and (size_val is None or size_val > 0):
             raise ValueError(f"Out of bound slice: start={value}, dim={s}.")
-          if size := _maybe_concretize(idx.size):
-            if value + (size - 1) * idx.stride >= s:
+          if size_val:
+            if value + (size_val - 1) * idx.stride >= s:
               raise ValueError(
-                  f"Out of bound slice: start={value}, size={size},"
+                  f"Out of bound slice: start={value}, size={size_val},"
                   f" stride={idx.stride}, dim={s}."
               )
         continue

jax/_src/state/types.py

@@ -271,18 +271,110 @@ def pretty_print(self, context: core.JaxprPpContext) -> pp.Doc:
     return pp.text(f"{{{self}}}")
 
 
+@tree_util.register_pytree_node_class
+@dataclasses.dataclass(frozen=True)
+class RefNewAxis:
+  """Transform that inserts new axes at specified positions."""
+  positions: tuple[int, ...]  # positions to insert new axes (in output)
+
+  def tree_flatten(self):
+    return (), (self.positions,)
+
+  @classmethod
+  def tree_unflatten(cls, metadata, arrays):
+    assert not arrays
+    return cls(*metadata)
+
+  def transform_shape(
+      self, shape: tuple[int | Array, ...] | None
+  ) -> tuple[int | Array, ...] | None:
+    if shape is None:
+      return None
+    result = list(shape)
+    for pos in sorted(self.positions):
+      result.insert(pos, 1)
+    return tuple(result)
+
+  def transform_dtype(self, dtype):
+    return dtype
+
+  def transform_sharding(self, sharding):
+    if all(p is None for p in sharding.spec):
+      return sharding
+    raise NotImplementedError
+
+  def pretty_print(self, context: core.JaxprPpContext) -> pp.Doc:
+    del context
+    return pp.text(f"{{newaxis{list(self.positions)}}}")
+
+
 @dataclasses.dataclass
 class RefIndexer:
   ref_or_view: Any
 
   def __getitem__(self, slc):
     if not isinstance(slc, tuple):
       slc = (slc,)
-    indexer = indexing.NDIndexer.from_indices_shape(slc, self.ref_or_view.shape)
+
+    # Handle None values (np.newaxis) in indices
+    none_positions = []
+    filtered_indices = []
+    output_pos = 0
+
+    # First pass: expand ellipsis if present
+    indices = list(slc)
+    num_none = sum(idx is None for idx in indices)
+    num_ellipsis = sum(idx is ... for idx in indices)
+
+    if num_ellipsis > 0:
+      # Expand ellipsis accounting for None values
+      ip = indices.index(...)
+      num_real_indices = len(indices) - num_ellipsis - num_none
+      num_slices_needed = len(self.ref_or_view.shape) - num_real_indices
+      indices[ip:ip+1] = [slice(None)] * max(0, num_slices_needed)
+
+    # Second pass: separate None from other indices and track positions
+    for idx in indices:
+      if idx is None:
+        none_positions.append(output_pos)
+        output_pos += 1
+      else:
+        filtered_indices.append(idx)
+        # Slices and ints both consume one input dim
+        # but slices produce one output dim, ints produce zero
+        if isinstance(idx, slice) or isinstance(idx, indexing.Slice):
+          output_pos += 1
+        elif not isinstance(idx, (int, np.integer)) and hasattr(idx, 'shape') and idx.shape:
+          # Array indexer produces its shape
+          output_pos += len(idx.shape)
+        # Scalar int indexers don't add to output dims
+
+    # Create indexer without None
+    filtered_tuple = tuple(filtered_indices) if filtered_indices else (slice(None),) * len(self.ref_or_view.shape)
+    if not filtered_indices and self.ref_or_view.shape == ():
+      # Special case: scalar ref with only None indexing
+      filtered_tuple = ()
+
+    # Build the result
     if isinstance(self.ref_or_view, TransformedRef):
       view = self.ref_or_view
-      return TransformedRef(view.ref, (*view.transforms, indexer))
-    return TransformedRef(self.ref_or_view, (indexer,))
+      if filtered_tuple:
+        indexer = indexing.NDIndexer.from_indices_shape(filtered_tuple, view.shape)
+        transforms = (*view.transforms, indexer)
+      else:
+        transforms = view.transforms
+      if none_positions:
+        transforms = (*transforms, RefNewAxis(tuple(none_positions)))
+      return TransformedRef(view.ref, transforms)
+    else:
+      if filtered_tuple:
+        indexer = indexing.NDIndexer.from_indices_shape(filtered_tuple, self.ref_or_view.shape)
+        transforms = (indexer,)
+      else:
+        transforms = ()
+      if none_positions:
+        transforms = (*transforms, RefNewAxis(tuple(none_positions)))
+      return TransformedRef(self.ref_or_view, transforms)
 
 
 @dataclasses.dataclass(frozen=True)