Implement hermitian conjugate utilities

jax/_src/numpy/array_methods.py

@@ -106,6 +106,20 @@ def _transpose(a: Array, *args: Any) -> Array:
   return lax_numpy.transpose(a, axis)
 
 
+def _hermitian(a: Array, *args: Any) -> Array:
+  """Returns a view of the array with axes transposed and values conjugated.
+
+  Refer to :func:`jax.numpy.hermitian` for full documentation.
+  """
+  if not args:
+    axis = None
+  elif len(args) == 1:
+    axis = args[0] if args[0] is None else _ensure_index_tuple(args[0])
+  else:
+    axis = _ensure_index_tuple(args)
+  return lax_numpy.hermitian(a, axis)
+
+
 def _compute_newshape(a: ArrayLike, newshape: DimSize | Shape) -> Shape:
   """Fixes a -1 value in newshape, if present."""
   orig_newshape = newshape  # for error messages
@@ -673,6 +687,7 @@ def max(self, values, *, indices_are_sorted=False, unique_indices=False,
   "diagonal": lax_numpy.diagonal,
   "dot": lax_numpy.dot,
   "flatten": lax_numpy.ravel,
+  "hermitian": _hermitian,
   "item": _item,
   "max": reductions.max,
   "mean": reductions.mean,
@@ -710,6 +725,8 @@ def max(self, values, *, indices_are_sorted=False, unique_indices=False,
   "flat": _notimplemented_flat,
   "T": lax_numpy.transpose,
   "mT": lax_numpy.matrix_transpose,
+  "H": lax_numpy.hermitian,
+  "mH": lax_numpy.matrix_hermitian,
   "real": ufuncs.real,
   "imag": ufuncs.imag,
   "nbytes": _nbytes,

jax/_src/numpy/lax_numpy.py

@@ -564,6 +564,51 @@ def matrix_transpose(x: ArrayLike, /) -> Array:
   return lax.transpose(x, axes)
 
 
+def hermitian(a: ArrayLike, axes: Sequence[int] | None = None) -> Array:
+  """Returns the hermitian conjugate transpose of an array.
+
+  Parameters
+  ----------
+  a : array_like
+      Input array.
+  axes : tuple or list of ints, optional
+      If specified, it must be a tuple or list which contains a permutation of [0,1,…,
+      N-1] where N is the number of axes of a. The i’th axis of the returned array will
+      correspond to the axis numbered `axes[i]` of the input. If not specified,
+      defaults to `range(a.ndim)[::-1]`, which reverses the order of the axes.
+
+  Returns
+  -------
+  out : array_like
+      The hermitian conjugate transpose of `a`.
+  """
+  util.check_arraylike("hermitian", a)
+  axes_ = list(range(ndim(a))[::-1]) if axes is None else axes
+  axes_ = [_canonicalize_axis(i, ndim(a)) for i in axes_]
+  return lax.conj(lax.transpose(a, axes_))
+
+
+def matrix_hermitian(x: ArrayLike, /) -> Array:
+  """Hermitian conjugate the last two dimensions of x.
+
+  Parameters
+  ----------
+  x : array_like
+      Input array. Must have ``x.ndim >= 2``.
+
+  Returns
+  -------
+  xH : Array
+      Hermitian conjugated array.
+  """
+  util.check_arraylike("matrix_hermitian", x)
+  ndim = np.ndim(x)
+  if ndim < 2:
+    raise ValueError(f"x must be at least two-dimensional for matrix_transpose; got {ndim=}")
+  axes = (*range(ndim - 2), ndim - 1, ndim - 2)
+  return lax.conj(lax.transpose(x, axes))
+
+
 @util._wraps(np.rot90, lax_description=_ARRAY_VIEW_DOC)
 @partial(jit, static_argnames=('k', 'axes'))
 def rot90(m: ArrayLike, k: int = 1, axes: tuple[int, int] = (0, 1)) -> Array: