Static typing for named tensor dimensions

[j-emberton]

Feature idea: Static typing for named tensor dimensions (xarray, Coordax, Bilt)

I've been looking at Pyrefly's new [tensor shape typing support](https://pyrefly.org/en/docs/tensor-shapes/), and wondered whether it might be possible to extend this to support named tensor dimensions, as used by libraries such as [xarray](https://github.com/pydata/xarray), [Coordax](https://github.com/neuralgcm/coordax), and [Bilt](https://github.com/aviation-impact-accelerator/bilt).

Pyrefly already supports symbolic shapes and shape transformations, so it seems there might be an opportunity to build on this infrastructure.

Motivation

Several scientific Python libraries use named dimensions rather than identifying tensor axes solely by their position and size:

• xarray provides labelled multidimensional arrays with broadcasting and alignment based on dimension names.
• Coordax brings labelled axes to JAX, supporting JAX transformations and interoperability with xarray. It was developed for applications such as NeuralGCM.
• Bilt uses named tensor domains in an equation-based modelling system, where equations can operate over different dimensions and propagate domains through calculations.

For example, two tensors might both have shape (100,), but one represents time and the other location. These dimensions are semantically different, despite having identical shapes.

Named dimensions also change broadcasting behaviour. Conceptually:

a: Tensor[Time, Location]
b: Tensor[Location]

c = a * b  # Tensor[Time, Location]

And:

a: Tensor[Time]
b: Tensor[Location]

c = a * b  # Tensor[Time, Location]

Unlike conventional positional broadcasting, the result's dimensions are determined by the union of the input dimensions.

Similarly, reducing over a named dimension removes it from the result:

a: Tensor[Time, Location]

b = a.mean(over=Time)  # Tensor[Location]

Ideally, a static checker could infer these domains and detect mismatches between inferred and declared tensor types.

Possible approach

I wondered whether Pyrefly's existing shape transformation DSL could be extended to support named dimensions and a small set of domain operations, such as:

• domain_union(A, B) — combine named dimensions for broadcasting.
• domain_remove(A, D) — remove a dimension during reduction.
• domain_equal(A, B) — compare domains independently of dimension ordering.
• domain_contains(A, D) — check membership of a named dimension.

Conceptually, this might allow library authors to define typing behaviour along the lines of:

def __mul__[A: Domain, B: Domain](
    self: Tensor[A],
    other: Tensor[B],
) -> Tensor[domain_union(A, B)]:
    ...

This is illustrative syntax rather than a proposal for a specific implementation.

A more ambitious approach might allow named dimensions to carry symbolic sizes:

Tensor[Dim["time", T], Dim["location", L]]

This could support both conventional shape checking and named-dimension compatibility.

Questions

I'm interested in whether:

1. The existing shape typing infrastructure could be generalised to support named dimensions, or whether this would require substantial changes to Pyrefly's type inference engine.
2. Domain operations could be implemented through the existing DSL, rather than requiring library-specific logic in Pyrefly.
3. There is a way to experiment with this using the current implementation, perhaps through symbolic dimension identities or custom stubs.
4. This could be implemented as a general capability usable by libraries such as xarray, Coordax and Bilt, each of which has somewhat different runtime semantics.

For context, we're exploring stronger static typing for Bilt. Being able to validate domain propagation and equation return types statically would be particularly valuable.

Given that Coordax already brings labelled dimensions into JAX-based scientific computing, this feels like an area where static typing could be useful beyond any one library.

I'd be interested to hear whether this aligns with the direction of Pyrefly's tensor typing work, and whether there might be a sensible starting point for a prototype.

[kinto0]

cc @stroxler

[stroxler]

Hi @j-emberton!

At the moment, I'm not sure how best to handle named arrays. I would very much like to support all of these libraries, and I think we should be able to do it. The tricky thing is figuring out how best to combine this with the type system we already have.

I'm confident we can find ways to do it, but I would want to take time to think about how best to handle it, especially at the boundary between shapes and names (like if you pull a plain JAX array out of a Coordax array, making sure Pyrefly can give the best available behavior).

When I think about this, it also reminds me of distributed torch and JAX sharding, where the axes of the arrays per se aren't named but there's a separate mesh grid with named axes that can be associated. I hope we can find a solution space flexible enough to support both use cases (cc @jakevdp - you mentioned someone on JAX was asking about sharding and shape types, this is a good conversation to pull them in!).

The biggest-picture question I have is the fundamental type layout: should we actually model this sort of thing as multiple aligned tuples of data (which is the natural thing to do), or should we model axes more flexibly so that it's possible to describe libraries / specific use cases (like sharding) where the axes can carry various kinds of metadata.

This question (of how to best model metadata per axis) actually impacts not just named arrays, but also kernel authoring where the strides are often important as well. The prototypes we've done so far for kernel authoring model shape and strides as independent data structures, but I keep wondering whether it would be better to go in an axis-as-metadata direction (cc @avikchaudhuri - I keep forgetting to mention this nagging feeling I have that treating shape and stride separately may not be optimal).

At some point I can write up a proposal but I would need time to think about this and I'm not confident I understand what the right ergonomics would be from the library / user side because I've never worked with sharding or named arrays myself (I used to do a lot of numpy so I have more confidence in my instincts there). But also, if anyone on this thread wants to take a stab at proposing how the type system could work I'd love to see suggestions informed by library author and user experience.

[j-emberton]

Thanks @stroxler, this is really interesting!

A consideration from Bilt is that named dimensions are independent of axis ordering. For example, Tensor[Time, Location] and Tensor[Location, Time] represent the same logical domain, and broadcasting operates on the union of named dimensions rather than their positions.

I think it would be useful to distinguish between three things:

Logical dimensions — their identity and size.

Physical axes — their position, shape and potentially strides.

Mappings between them — including how axes relate to named dimensions or sharding meshes.

I'd be happy to help put together some examples from Bilt to explore how this might work.