Is it practical to use axis names in runtime functionality?

[ships]

Please forgive my youth as I wrap my head around the constraints here, and whether what I am trying to do has some sanity to it <3
I have a PyTree (actually, eqx module, actually a pair of them) like this:

class A(eqx.Module):
    p: Integer[Array, "*batches xs foos"]
    q: Integer[Array, "*batches xs bazs"]

class B(eqx.Module):
    m: Integer[Array, ""]
    r: Integer[Array, "*batches xs foos bangs"]
    s: Integer[Array, "*batches xs bazs bangs"]

and so on.

My desire is to write functions that operate reliably on pytrees that have several leaves with shared axis dims, modifying those dims, preferably by name. I can write variants for each function that is type-gnostic, but that is toil I hope to learn to avoid -- because if I understand correctly, in principle all the info is in the name!

So, a trivial example function, written in the most redundant way:

@dispatch
def nth_x(it: A, n: int) -> A:
    return A(
        p=it.p[..., n, :],
        q=it.q[..., n, :],
    )

@dispatch
def nth_x(it: B, n: int) -> B:
    return B(
        m=it.m,
        r=it.r[..., n, :, :],
        s=it.s[..., n, :, :],
    )


@dispatch
def nth_baz(it: A, n: int) -> A:
    return A(
        p=it.p, # reduction from `it.p[:,:,:]` because bazs is not one of the dims
        q=it.q[..., :, n],
    )

@dispatch
def nth_baz(it: B, n: int) -> B:
    return B(
        m=it.m,
        r=it.r, # reduction from `it.r[:,:,:,:]` because bazs is not one of the dims
        s=it.s[..., :, n, :],
    )

Note that the indexing differs per leaf — xs is axis -2 in A's leaves but axis -3 in B's r and s, and absent entirely from m. Each variant is trivial to write, but I'd need one per type, per operation. Less redundant would be something like this:

def nth_x(it: PyTree, n: int) -> PyTree:
    return named_index(it, {"xs": n})

def swap_first_last_bang(it: PyTree) -> PyTree:
    return named_index(it, {"bangs": slice(None, None, -1)})

# or even composing multiple axes at once:
def nth_x_first_few_foos(it: PyTree, n: int, k: int) -> PyTree:
    return named_index(it, {"xs": n, "foos": slice(0, k)})

This is obviously nontrivial for many reasons and i have a hunch it is not possible without, maybe, some bonkers use of the debug api. And the dict-like setup is just trying to reconcile names with nice array indexing syntax, but is hideous. OTOH, it occurs to me i might just be looking at the elephant from the wrong side and you have a tip about the right way to approach this. FWIW I have been able to skirt around this in a few cases because by coincidence of the setting, i.e. with some assumptions where all arrays are either scalar or contain the vary axis of interest. But i have a case currently where that is finally strictly untrue, i.e. as shown above, there are several NDarrays, some shared batch axes, but 1 array is missing a cardinal dimension and is semantically "static [or broadcast?] relative to this batch dim/operation".

Do you have any guidance? Thank you very much, this project (and equinox) are such big helps.

[patrick-kidger]

So you can access the type annotation with e.g. inspect.get_annotations(A), and given an annotation ann = Float[...] then you can access the original dimension string with ann.dim_str or the processed dimensions with ann.dims. You could probably combine those to set up your logic.

That said, this is definitely not public API so I offer no guarantees about breaking it!

As an alternative strategy you might like Haliax, which is designed to offer named arrays with Equinox.

[ships]

Thanks for the thoughts. I ran into Haliax, and it looks like it could cover this :) but is a larger lift to rewrite stuff than i can do imminently. will write back here once i get a chance to explore it. Thank you!