more functions

Author: jakevdp

jax/experimental/array_api/__init__.py

@@ -127,3 +127,18 @@
     tanh as tanh,
     trunc as trunc,
 )
+
+from jax.experimental.array_api._indexing_functions import take as take
+
+from jax.experimental.array_api._manipulation_functions import (
+    broadcast_arrays as broadcast_arrays,
+    broadcast_to as broadcast_to,
+    concat as concat,
+    expand_dims as expand_dims,
+    flip as flip,
+    permute_dims as permute_dims,
+    reshape as reshape,
+    roll as roll,
+    squeeze as squeeze,
+    stack as stack,
+)

jax/experimental/array_api/_indexing_functions.py

@@ -0,0 +1,18 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax
+
+def take(x, indices, /, *, axis):
+  return jax.numpy.take(x, indices, axis=axis)

jax/experimental/array_api/_manipulation_functions.py

@@ -0,0 +1,71 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax
+from jax.experimental.array_api._data_type_functions import result_type as _result_type
+
+
+def broadcast_arrays(*arrays):
+  """Broadcasts one or more arrays against one another."""
+  return jax.numpy.broadcast_arrays(*arrays)
+
+
+def broadcast_to(x, /, shape):
+  """Broadcasts an array to a specified shape."""
+  return jax.numpy.broadcast_to(x, shape=shape)
+
+
+def concat(arrays, /, *, axis=0):
+  """Joins a sequence of arrays along an existing axis."""
+  dtype = _result_type(*arrays)
+  return jax.numpy.concatenate(arrays, axis=axis, dtype=dtype)
+
+
+def expand_dims(x, /, *, axis=0):
+  """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):
+  """Reverses the order of elements in an array along the given axis."""
+  dimensions = list(axis) if isinstance(axis, tuple) else [axis]
+  return jax.lax.rev(x, dimensions=dimensions)
+
+
+def permute_dims(x, /, axes):
+  """Permutes the axes (dimensions) of an array x."""
+  return jax.lax.transpose(x, axes)
+
+
+def reshape(x, /, shape, *, copy=None):
+  """Reshapes an array without changing its data."""
+  del copy  # unused
+  return jax.lax.reshape(x, shape)
+
+
+def roll(x, /, shift, *, axis=None):
+  """Rolls array elements along a specified axis."""
+  return jax.numpy.roll(x, shift=shift, axis=axis)
+
+
+def squeeze(x, /, axis):
+  """Removes singleton dimensions (axes) from x."""
+  dimensions = list(axis) if isinstance(axis, tuple) else [axis]
+  return jax.lax.squeeze(x, dimensions=dimensions)
+
+
+def stack(arrays, /, *, axis=0):
+  """Joins a sequence of arrays along a new axis."""
+  dtype = _result_type(*arrays)
+  return jax.numpy.stack(arrays, axis=axis, dtype=dtype)