more functions

Author: jakevdp

jax/experimental/array_api/__init__.py

@@ -173,3 +173,15 @@
     sum as sum,
     var as var
 )
+
+from jax.experimental.array_api._utility_functions import (
+    all as all,
+    any as any,
+)
+
+from jax.experimental.array_api._linear_algebra_functions import (
+    matmul as matmul,
+    matrix_transpose as matrix_transpose,
+    tensordot as tensordot,
+    vecdot as vecdot,
+)

jax/experimental/array_api/_linear_algebra_functions.py

@@ -0,0 +1,65 @@
+# 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.
+
+from typing import NamedTuple
+
+import jax
+
+
+class EighResult(NamedTuple):
+  eigenvalues: jax.Array
+  eigenvectors: jax.Array
+
+class QRResult(NamedTuple):
+  Q: jax.Array
+  R: jax.Array
+
+class SlogdetResult(NamedTuple):
+  sign: jax.Array
+  logabsdet: jax.Array
+
+class SVDResult(NamedTuple):
+  U: jax.Array
+  S: jax.Array
+  Vh: jax.Array
+
+
+def matmul(x1, x2, /):
+  """Computes the matrix product."""
+  return jax.numpy.matmul(x1, x2)
+
+
+def matrix_transpose(x, /):
+  """Transposes a matrix (or a stack of matrices) x."""
+  if x.ndim < 2:
+    raise ValueError(f"matrix_transpose requres at least 2 dimensions; got {x.ndim=}")
+  return jax.lax.transpose(x, (*range(x.ndim - 2), x.ndim - 1, x.ndim - 2))
+
+
+def tensordot(x1, x2, /, *, axes=2):
+  """Returns a tensor contraction of x1 and x2 over specific axes."""
+  return jax.numpy.tensordot(x1, x2, axes=axes)
+
+
+def vecdot(x1, x2, /, *, axis=-1):
+  """Computes the (vector) dot product of two arrays."""
+  rank = max(x1.ndim, x2.ndim)
+  x1 = jax.lax.broadcst_to_rank(x1, rank)
+  x2 = jax.lax.broadcast_to_rank(x2, rank)
+  if x1.shape[axis] != x2.shape[axis]:
+    raise ValueError("x1 and x2 must have the same size along specified axis.")
+  x1, x2 = jax.lax.broadcast_arrays(x1, x2)
+  x1 = jax.numpy.moveaxis(x1, axis, -1)
+  x2 = jax.numpy.moveaxis(x2, axis, -1)
+  return jax.numpy.matmul(x1[..., None, :], x2[..., None])[..., 0, 0]

jax/experimental/array_api/_sorting_functions.py

@@ -28,4 +28,3 @@ def sort(x, /, *, axis=-1, descending=False, stable=True):
   if descending:
     return jax.lax.rev(result, dimensions=[axis])
   return result
-

jax/experimental/array_api/_utility_functions.py

@@ -0,0 +1,25 @@
+# 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 all(x, /, *, axis=None, keepdims=False):
+  """Tests whether all input array elements evaluate to True along a specified axis."""
+  return jax.numpy.all(x, axis=axis, keepdims=keepdims)
+
+
+def any(x, /, *, axis=None, keepdims=False):
+  """Tests whether any input array element evaluates to True along a specified axis."""
+  return jax.numpy.any(x, axis=axis, keepdims=keepdims)