Alternative data model for NamedArray using tuple[str | None, ...] for axis names

[shoyer]

I've implemented a fork of Penzai using tuple[str | None, ...] for axis names:
https://github.com/neuralgcm/neuralgcm/blob/715b2d6c6a4d2bb08c4cbdeb20d28ce43f957471/neuralgcm/experimental/coordax/named_axes.py

I like this model a little better, because it's simpler to understand and only includes a single NamedArray type.

This works with JAX's tree operations to handle operations like scan/vmap by allowing for inserting or trimming unnamed leading dimensions:
https://github.com/neuralgcm/neuralgcm/blob/715b2d6c6a4d2bb08c4cbdeb20d28ce43f957471/neuralgcm/experimental/coordax/named_axes.py#L443

[danieldjohnson]

Very cool! I'll dig into this more when I have time.

It looks like you've made it so named axis order matters? So arrays with names in different orders are distinct from the perspective of the API. That seems like a pretty big difference in mental model, I wonder how much that would affect how it feels to use. (I suppose you can always reorder right before you untag using order_as?)

Also, I'm always a bit hesitant about types that have custom PyTree flattening / unflattening logic. Some constraints I noticed:

• It looks like you don't allow changing the size of named axes inside PyTree flattening / unflattening ops?
• Also, I think this implementation supports (NamedArray with arrays -> flatten -> replace with non-array data-> unflatten) but not (NamedArray with non-array data -> flatten -> add arrays -> unflatten).
  I wonder if this breaks anything. It seems like it might break code that is assuming PyTree nodes are just containers.

[shoyer]

It looks like you've made it so named axis order matters? So arrays with names in different orders are distinct from the perspective of the API.

Yes, indeed. But so far I've avoided adding APIs to NamedArray itself where this would actually made a difference, aside from the .data property itself.

My general strategy when building Coordax's NamedArray has been to focus the API around the case where axis is named, but to add enough flexibility to support unnamed axes for interfacing with other code (via nmap) and JAX transformation. The recommended workflow will be to always fully name axes after interfacing with other JAX code, so methods like order_as work.

It looks like you don't allow changing the size of named axes inside PyTree flattening / unflattening ops?

Yes, this is an attempt to match Penzai's behavior. I think verifying that named axes haven't changed since is probably a useful constraint to catch bugs, but we'll have to see how helpful that is in practice.

Also, I think this implementation supports (NamedArray with arrays -> flatten -> replace with non-array data-> unflatten) but not (NamedArray with non-array data -> flatten -> add arrays -> unflatten).
I wonder if this breaks anything. It seems like it might break code that is assuming PyTree nodes are just containers.

I think it is probably safe to assume that non-arrays will only get inserted into flattened pytrees, because the pytree interface makes no guarantees about constructors. But I share your concern about fragility of my implementation. This will be something worth checking with the JAX team. I have mostly been relying on testing with JAX transformations to verify it works as expectation, rather than following documented stability guarantees.