Merge pull request #24932 from hawkinsp:gather

PiperOrigin-RevId: 697632284
jax-ml · Nov 18, 2024 · 65f9c78 · 65f9c78
2 parents ccb3317 + 1d519f4
commit 65f9c78
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
@@ -4544,6 +4544,8 @@ def shape_as_value(shape: core.Shape):
   """Converts a shape that may contain Poly values into a JAX value."""
   if len(shape) == 0:
     return full((0,), np.array(0, np.int64))
+  if core.is_constant_shape(shape):
+    return np.asarray(shape, dtype=np.int64)
   dims = [
       expand_dims(convert_element_type(core.dimension_as_value(d), np.int64),
                   (0,))