jnp.take_along_axis: support fill_value

Author: jakevdp

jax/_src/numpy/lax_numpy.py

@@ -5725,12 +5725,13 @@ def _normalize_index(index, axis_size):
 
 @util.implements(np.take_along_axis, update_doc=False,
         lax_description=TAKE_ALONG_AXIS_DOC)
-@partial(jit, static_argnames=('axis', 'mode'))
+@partial(jit, static_argnames=('axis', 'mode', 'fill_value'))
 def take_along_axis(
     arr: ArrayLike,
     indices: ArrayLike,
     axis: int | None,
     mode: str | lax.GatherScatterMode | None = None,
+    fill_value: StaticScalar | None = None,
 ) -> Array:
   util.check_arraylike("take_along_axis", arr, indices)
   a = asarray(arr)
@@ -5743,8 +5744,9 @@ def take_along_axis(
     if ndim(indices) != 1:
       msg = "take_along_axis indices must be 1D if axis=None, got shape {}"
       raise ValueError(msg.format(idx_shape))
-    return take_along_axis(a.ravel(), indices, 0)
-  rank = ndim(arr)
+    a = a.ravel()
+    axis = 0
+  rank = a.ndim
   if rank != ndim(indices):
     msg = "indices and arr must have the same number of dimensions; {} vs. {}"
     raise ValueError(msg.format(ndim(indices), a.ndim))
@@ -5812,7 +5814,7 @@ def replace(tup, val):
     collapsed_slice_dims=tuple(collapsed_slice_dims),
     start_index_map=tuple(start_index_map))
   return lax.gather(a, gather_indices_arr, dnums, tuple(slice_sizes),
-                    mode="fill" if mode is None else mode)
+                    mode="fill" if mode is None else mode, fill_value=fill_value)
 
 
 ### Indexing

jax/numpy/__init__.pyi

@@ -811,6 +811,7 @@ def take_along_axis(
     indices: ArrayLike,
     axis: Optional[int],
     mode: Optional[Union[str, GatherScatterMode]] = ...,
+    fill_value: Optional[StaticScalar] = None,
 ) -> Array: ...
 def tan(x: ArrayLike, /) -> Array: ...
 def tanh(x: ArrayLike, /) -> Array: ...

tests/lax_numpy_test.py

@@ -4483,6 +4483,13 @@ def testTakeAlongAxisWithEmptyArgs(self):
     x = jnp.ones((4, 0, 3), dtype=jnp.int32)
     np.testing.assert_array_equal(x, jnp.take_along_axis(x, x, axis=1))
 
+  def testTakeAlongAxisOptionalArgs(self):
+    x = jnp.arange(5.0)
+    ind = jnp.array([0, 2, 4, 6])
+    expected = jnp.array([0.0, 2.0, 4.0, 10.0], dtype=x.dtype)
+    actual = jnp.take_along_axis(x, ind, axis=None, mode='fill', fill_value=10.0)
+    self.assertArraysEqual(expected, actual)
+
   @jtu.sample_product(
     dtype=inexact_dtypes,
     shape=[0, 5],
@@ -5973,7 +5980,7 @@ def testWrappedSignaturesMatch(self):
       'clip': ['x', 'max', 'min'],
       'einsum': ['subscripts', 'precision'],
       'einsum_path': ['subscripts'],
-      'take_along_axis': ['mode'],
+      'take_along_axis': ['mode', 'fill_value'],
       'fill_diagonal': ['inplace'],
     }