def _nmap_with_doc(
    fun: Callable[..., Any], fun_name: str, fun_doc: str | None = None
) -> Callable[..., Any]:
  """Builds a nmap-wrapped function with a docstring."""

  @functools.wraps(fun)
  def wrapped_fun(*args, **kwargs):
    arg_leaves_and_paths, arg_treedef = jax.tree_util.tree_flatten_with_path(
        (args, kwargs),
        is_leaf=lambda node: isinstance(node, NamedArray | NamedArrayView),
    )
    arg_leaves = [leaf for _, leaf in arg_leaves_and_paths]
    # Extract any argument leaves that were NamedArrays or NamedArrayViews. The
    # rest of the arguments will just be closed over, so they don't have to be
    # arraylike. We also check named shapes here.
    # To simplify the implementation, we ensure that all arguments are views
    # rather than positional-prefix NamedArrays; this can always be done without
    # any device computations.
    named_array_arg_leaves = []
    known_sizes = {}
    bad_names = []
    for leaf in arg_leaves:
      if isinstance(leaf, NamedArray | NamedArrayView):
        for name, size in leaf.named_shape.items():
          if name in known_sizes:
            if known_sizes[name] != size and name not in bad_names:
              bad_names.append(name)
          else:
            known_sizes[name] = size
        named_array_arg_leaves.append(leaf.as_namedarrayview())

    if bad_names:
      msg = [
          f"Inconsistent named axes in a call to nmap({fun}) for axes"
          f" {bad_names}:"
      ]
      for keypath, leaf in arg_leaves_and_paths:
        if isinstance(leaf, NamedArray | NamedArrayView):
          assert keypath
          if keypath[0] == jax.tree_util.SequenceKey(0):
            prefix = f"args{jax.tree_util.keystr(keypath[1:])}"
          elif keypath[0] == jax.tree_util.SequenceKey(1):
            prefix = f"kwargs{jax.tree_util.keystr(keypath[1:])}"
          else:
            # Should never happen!
            prefix = f"tree{jax.tree_util.keystr(keypath)}"
          msg.append(f"  {prefix}.named_shape == {leaf.named_shape}")
      raise ValueError("\n".join(msg))

    # Prepare a version of the function that accepts tracers for each of those
    # extracted arguments, and rebuilds the full arguments to `fun`.
    def flat_array_fun(batch_tracers):
      # Replace each NamedArray with its batch tracer.
      batch_tracers_stack = batch_tracers[::-1]
      new_arg_leaves = []
      for leaf in arg_leaves:
        if isinstance(leaf, NamedArray | NamedArrayView):
          # Substitute the tracer.
          new_arg_leaves.append(batch_tracers_stack.pop())
        else:
          new_arg_leaves.append(leaf)
      # Call the function.
      args, kwargs = jax.tree_util.tree_unflatten(arg_treedef, new_arg_leaves)
      return fun(*args, **kwargs)

    # Collect the axis names. This determines what we will be vectorizing over.
    # We use an ordered set to guarantee a deterministic ordering based on the
    # arguments, since sets have nondeterministic iteration order in Python.
    all_names = ordered_set.OrderedSet()
    for named_arg in named_array_arg_leaves:
      all_names.update(named_arg.data_axis_for_name.keys())

    all_names = list(all_names)

    # Recursively vectorize over each axis.
    def recursive_vectorize_step(current_views, remaining_names):
      if not remaining_names:
        # All names have been processed, so none of the args should have names.
        # Unwrap them and call the function.
        return flat_array_fun([view.unwrap() for view in current_views])

      # Otherwise, we still have names to vectorize over. Pop one name off the
      # stack and vectorize over it as needed.
      vmap_name = remaining_names[0]
      reduced_views = []
      vmap_axes = []
      for view in current_views:
        if vmap_name in view.data_axis_for_name:
          vmap_axis = view.data_axis_for_name[vmap_name]
          vmap_axes.append(vmap_axis)

          # pylint: disable=cell-var-from-loop
          def _shift_axis(other_axis):
            assert other_axis != vmap_axis
            if other_axis < vmap_axis:
              return other_axis
            else:
              return other_axis - 1

          # pylint: enable=cell-var-from-loop

          # We are temporarily constructing an "invalid" view here because
          # data_array will still have the extra axis. But after running `vmap`,
          # it will be valid again.
          reduced_views.append(
              NamedArrayView(
                  data_array=view.data_array,
                  data_axis_for_name={
                      name: _shift_axis(data_axis)
                      for name, data_axis in view.data_axis_for_name.items()
                      if name != vmap_name
                  },
                  data_axis_for_logical_axis=tuple(
                      _shift_axis(data_axis)
                      for data_axis in view.data_axis_for_logical_axis
                  ),
                  data_shape=(
                      view.data_shape[:vmap_axis]
                      + view.data_shape[vmap_axis + 1 :]
                  ),
              )
          )
        else:
          # This argument doesn't have this axis, so don't map over anything.
          vmap_axes.append(None)
          reduced_views.append(view)

      return jax.vmap(
          functools.partial(
              recursive_vectorize_step,
              remaining_names=remaining_names[1:],
          ),
          in_axes=(vmap_axes,),
          out_axes=0,
          axis_name=vmap_name,
      )(reduced_views)

    # Run the function.
    result_data = recursive_vectorize_step(named_array_arg_leaves, all_names)

    # Wrap all leaves in NamedArray or NamedArrayView, assigning the names from
    # `all_names` to their mapped-over axes. The mapped-over named axes always
    # end up at the front, followed by positional axes, so if there are any
    # positional axes we need to return a NamedArrayView.
    def handle_result(leaf):
      leaf = jnp.array(leaf)
      if leaf.ndim == len(all_names):
        return NamedArray(
            data_array=leaf,
            named_axes=collections.OrderedDict(zip(all_names, leaf.shape)),
        )
      else:
        assert leaf.ndim > len(all_names)
        return NamedArrayView(
            data_array=leaf,
            data_shape=leaf.shape,
            data_axis_for_name={name: i for i, name in enumerate(all_names)},
            data_axis_for_logical_axis=tuple(range(len(all_names), leaf.ndim)),
        )

    return jax.tree_util.tree_map(handle_result, result_data)

  docstr = (
      f"Name-vectorized version of `{fun_name}`. Takes similar arguments as"
      f" `{fun_name}` but accepts and returns NamedArrays (or NamedArrayViews)"
      " in place of regular arrays."
  )
  if fun_doc:
    docstr += f"\n\nOriginal documentation:\n\n{fun_doc}"
  wrapped_fun.__doc__ = docstr
  return wrapped_fun
Builds a nmap-wrapped function with a docstring.
penzai/core/named_axes.py