register namespace

Author: jakevdp

jax/experimental/array_api/__init__.py

@@ -185,3 +185,18 @@
     tensordot as tensordot,
     vecdot as vecdot,
 )
+
+def _array_namespace(self, /, *, api_version: None | str = None):
+  import sys
+  if api_version is not None and api_version != __array_api_version__:
+    raise ValueError(f"{api_version=!r} is not available; "
+                     f"available versions are: {[__array_api_version__]}")
+  return sys.modules[__name__]
+
+def _setup_array_type():
+  # TODO(jakevdp): set on tracers as well?
+  from jax._src.array import ArrayImpl
+  setattr(ArrayImpl, "__array_namespace__", _array_namespace)
+
+_setup_array_type()
+del _setup_array_type

jax/experimental/array_api/_data_type_functions.py

@@ -122,10 +122,12 @@ def _promote_types(t1, t2):
 
 
 def astype(x, dtype, /, *, copy=True):
-  return jnp.asarray(x, dtype=dtype, copy=copy)
+  return jnp.array(x, dtype=dtype, copy=copy)
 
 
 def can_cast(from_, to, /):
+  if isinstance(from_, jax.Array):
+    from_ = from_.dtype
   if not _is_valid_dtype(from_):
     raise ValueError(f"{from_} is not a valid dtype")
   if not _is_valid_dtype(to):

jax/experimental/array_api/_manipulation_functions.py

@@ -12,60 +12,74 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import operator
+from typing import List, Optional, Tuple, Union
+
 import jax
+from jax import Array
 from jax.experimental.array_api._data_type_functions import result_type as _result_type
 
 
-def broadcast_arrays(*arrays):
+def broadcast_arrays(*arrays: Array) -> List[Array]:
   """Broadcasts one or more arrays against one another."""
   return jax.numpy.broadcast_arrays(*arrays)
 
 
-def broadcast_to(x, /, shape):
+def broadcast_to(x: Array, /, shape: Tuple[int]) -> Array:
   """Broadcasts an array to a specified shape."""
   return jax.numpy.broadcast_to(x, shape=shape)
 
 
-def concat(arrays, /, *, axis=0):
+def concat(arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: Optional[int] = 0) -> Array:
   """Joins a sequence of arrays along an existing axis."""
   dtype = _result_type(*arrays)
+  if axis is None:
+    arrays = [reshape(arr, (arr.size,)) for arr in arrays]
+    axis = 0
   return jax.numpy.concatenate(arrays, axis=axis, dtype=dtype)
 
 
-def expand_dims(x, /, *, axis=0):
+def expand_dims(x: Array, /, *, axis: int = 0) -> Array:
   """Expands the shape of an array by inserting a new axis (dimension) of size one at the position specified by axis."""
   return jax.lax.expand_dims(x, dimensions=[axis])
 
 
-def flip(x, /, *, axis=None):
+def flip(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> Array:
   """Reverses the order of elements in an array along the given axis."""
-  dimensions = list(axis) if isinstance(axis, tuple) else [axis]
+  if axis is None:
+    dimensions = tuple(range(x.ndim))
+  elif isinstance(axis, int):
+    dimensions = (axis,)
+  elif isinstance(axis, tuple):
+    dimensions = tuple(operator.index(ax) for ax in axis)
+  else:
+    raise TypeError(f"Unexpected input axis={axis}: expected None, int, or tuple of ints")
   return jax.lax.rev(x, dimensions=dimensions)
 
 
-def permute_dims(x, /, axes):
+def permute_dims(x: Array, /, axes: Tuple[int, ...]) -> Array:
   """Permutes the axes (dimensions) of an array x."""
   return jax.lax.transpose(x, axes)
 
 
-def reshape(x, /, shape, *, copy=None):
+def reshape(x: Array, /, shape: Tuple[int, ...], *, copy: Optional[bool] = None) -> Array:
   """Reshapes an array without changing its data."""
   del copy  # unused
   return jax.lax.reshape(x, shape)
 
 
-def roll(x, /, shift, *, axis=None):
+def roll(x: Array, /, shift: Union[int, Tuple[int]], *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> Array:
   """Rolls array elements along a specified axis."""
   return jax.numpy.roll(x, shift=shift, axis=axis)
 
 
-def squeeze(x, /, axis):
+def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array:
   """Removes singleton dimensions (axes) from x."""
-  dimensions = list(axis) if isinstance(axis, tuple) else [axis]
+  dimensions = axis if isinstance(axis, tuple) else (axis,)
   return jax.lax.squeeze(x, dimensions=dimensions)
 
 
-def stack(arrays, /, *, axis=0):
+def stack(arrays: Union[Tuple[Array, ...], List[Array]], /, *, axis: int = 0) -> Array:
   """Joins a sequence of arrays along a new axis."""
   dtype = _result_type(*arrays)
   return jax.numpy.stack(arrays, axis=axis, dtype=dtype)