API for surfacing active SPMD axis names

For example,

@partial(pmap, axis_name="i")
def foo():
  @partial(pmap, axis_name="j")
  def bar():
    print(lax.axis_names())  # ("i", "j")

The semantics would be “the set of axis names currently valid inside a psum or axis_index”; the order would be the pmap nesting order. cc @dmrd

[mattjj]

Regarding implementation, I think we should move the axis env to be part of the tracer state in core.py, for at least two reasons:

1. it centralizes the global state, which is useful when forming thunks, and
2. it’ll let vmap easily reuse the same axis environment.

[shoyer]

Another weird bit: this would be like functions being able to ask “how deep in the call stack am I?”

Yeah, this worries me a bit. This feels like a feature that might be good not to have, to ensure compositionality.

The batch normalization use case is real, though. I wonder if there could be a more hygenic way to do support this, e.g., only surfacing active axis names defined by a particular library?