def _wrap_array_method(name):
  """Wraps an array method on a named array."""

  def func(array, *args, **kwargs):
    return getattr(array, name)(*args, **kwargs)

  array_method = getattr(jax.Array, name)
  wrapped_func = nmap(func)
  functools.update_wrapper(
      wrapped_func,
      array_method,
      assigned=("__name__", "__qualname__", "__annotations__"),
      updated=(),
  )
  wrapped_func.__module__ = __name__
  wrapped_func.__doc__ = (
      "Name-vectorized version of array method"
      f" `{name} <numpy.ndarray.{name}>`. Takes similar arguments as"
      f" `{name} <numpy.ndarray.{name}>` but accepts and returns NamedArrays"
      " (or NamedArrayViews) in place of regular arrays."
  )
  return wrapped_func
Wraps an array method on a named array.
penzai/core/named_axes.py