Batched type

Hi, I was wondering if a type-annotation for Batched datastructures exists or if it could be implemented?

So something similar along the likes of PyTree but for indicating that all leaf nodes have a leading axis size of e.g., N.

I know I can already do this easily with e.g., PyTree[Float32[Array, 'N ...']], but in this way I can't use my type-aliases (or don't know how to).

Example use-case:

from jaxtyping import Array, Float32, PyTree
from jax import vmap

MyArray = Float32[Array, '...']
Batched = ???  # Should prepend 'N ' to Float32[Array, 'N ...']

def sample_fun(x: PyTree[MyArray]) -> PyTree[MyArray]:
    return x

def batch_fun(xs: Batched[PyTree[MyArray]]) -> PyTree[Batched[MyArray]]:
    return jax.vmap(sample_fun)(xs)

Note, intuitively Batched[PyTree[...]] should be equivalent to PyTree[Batched[...]] as they should operate on the leaves.

No implementation files or tests are named. Start by locating the existing PyTree annotation and type-alias handling, then determine how a Batched annotation should compose with both. Done means aliases can be wrapped to add the leading batch axis and Batched[PyTree[...]] behaves equivalently to PyTree[Batched[...]].