TypeVar compatibility with vmap or scan

[joeryjoery]

Hi I'm wondering if there is a way to type-annotate a TypeVar with a batch dimension.

Something that would make it compatible with vmap or scan of the sorts.

To give an example, I'm looking for something in place of Batched,

from typing import TypeVar
from jax import vmap

Batched = ???  # Does something like this exist?

T = TypeVar("T")

def sample_fun(x: T) -> T:
    return x

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

I have tried to do this with e.g., typing.Sequence, but my Type Checker is complaining that Array is not a Sequence.

Any hints would be greatly appreciated :)

[jakevdp]

There is no generally applicable annotation in this case, because xs may be any object compatible with pytree flattening, and this is not defined by any static type system but rather is defined by runtime pytree registration. So in the most general case, the best you can do is this:

def sample_fun(x: T) -> T:
    return x

def batch_fun(xs: T) -> T:
    return jax.vmap(sample_fun)(xs)

Of course, if you know that the input xs be a single array-like object rather than a general pytree-of-arrays, you could be more specific with something like this:

def sample_fun(x: Array) -> Array:
    return x

def batch_fun(xs: Array) -> Array:
    return jax.vmap(sample_fun)(xs)

From the static typing perspective, there is no way to express that an array is Batched – an array is an Array whether it's batched or not.