jnp.take: fix annotation for fill_value

Author: jakevdp

jax/_src/basearray.py

@@ -115,6 +115,14 @@ def sharding(self) -> Sharding:
 
 Array.__module__ = "jax"
 
+# StaticScalar is the Union of all scalar types that can be converted to
+# JAX arrays, and are possible to mark as static arguments.
+StaticScalar = Union[
+  np.bool_, np.number,  # NumPy scalar types
+  bool, int, float, complex,  # Python scalar types
+]
+StaticScalar.__doc__ = "Type annotation for JAX-compatible static scalars."
+
 
 # ArrayLike is a Union of all objects that can be implicitly converted to a
 # standard JAX array (i.e. not including future non-standard array types like
@@ -123,7 +131,6 @@ def sharding(self) -> Sharding:
 ArrayLike = Union[
   Array,  # JAX array type
   np.ndarray,  # NumPy array type
-  np.bool_, np.number,  # NumPy scalar types
-  bool, int, float, complex,  # Python scalar types
+  StaticScalar,  # valid scalars
 ]
 ArrayLike.__doc__ = "Type annotation for JAX array-like objects."

jax/_src/basearray.pyi

@@ -217,11 +217,15 @@ class Array(abc.ABC):
   def unsafe_buffer_pointer(self) -> int: ...
 
 
+StaticScalar = Union[
+  np.bool_, np.number,  # NumPy scalar types
+  bool, int, float, complex,  # Python scalar types
+]
+
 ArrayLike = Union[
   Array,  # JAX array type
   np.ndarray,  # NumPy array type
-  np.bool_, np.number,  # NumPy scalar types
-  bool, int, float, complex,  # Python scalar types
+  StaticScalar,  # valid scalars
 ]
 
 

jax/_src/numpy/lax_numpy.py

@@ -68,8 +68,8 @@
 from jax._src.numpy import util
 from jax._src.numpy.vectorize import vectorize
 from jax._src.typing import (
-  Array, ArrayLike, DimSize, DuckTypedArray,
-  DType, DTypeLike, Shape, DeprecatedArg
+  Array, ArrayLike, DeprecatedArg, DimSize, DuckTypedArray,
+  DType, DTypeLike, Shape, StaticScalar,
 )
 from jax._src.util import (unzip2, subvals, safe_zip,
                            ceil_of_ratio, partition_list,
@@ -5637,7 +5637,7 @@ def take(
     mode: str | None = None,
     unique_indices: bool = False,
     indices_are_sorted: bool = False,
-    fill_value: ArrayLike | None = None,
+    fill_value: StaticScalar | None = None,
 ) -> Array:
   return _take(a, indices, None if axis is None else operator.index(axis), out,
                mode, unique_indices=unique_indices, indices_are_sorted=indices_are_sorted,

jax/_src/typing.py

@@ -34,6 +34,7 @@
 from jax._src.basearray import (
     Array as Array,
     ArrayLike as ArrayLike,
+    StaticScalar as StaticScalar,
 )
 
 DType = np.dtype

jax/numpy/__init__.pyi

@@ -11,8 +11,8 @@ from jax._src.lax.slicing import GatherScatterMode
 from jax._src.lib import Device
 from jax._src.numpy.index_tricks import _Mgrid, _Ogrid, CClass as _CClass, RClass as _RClass
 from jax._src.typing import (
-    Array, ArrayLike, DType, DTypeLike,
-    DimSize, DuckTypedArray, Shape, DeprecatedArg
+    Array, ArrayLike, DType, DTypeLike, DeprecatedArg,
+    DimSize, DuckTypedArray, Shape, StaticScalar,
 )
 from jax.numpy import fft as fft, linalg as linalg
 from jax.sharding import Sharding as _Sharding
@@ -804,7 +804,7 @@ def take(
     mode: Optional[str] = ...,
     unique_indices: builtins.bool = ...,
     indices_are_sorted: builtins.bool = ...,
-    fill_value: Optional[ArrayLike] = ...,
+    fill_value: Optional[StaticScalar] = ...,
 ) -> Array: ...
 def take_along_axis(
     arr: ArrayLike,