Type hinting for `eqx.Module` created under `vmap`

Hello, I have something like the following pattern

import equinox as eqx
from jaxtyping import Array, Float

#
# Library code
#
class TestModule(eqx.Module):

    a: Float[Array, ""]
    b: Float[Array, "2"]

    def __init__(self, a, b):
        ...

def compute_something(vmapped_module: TestModule):
    """Specifically takes in TestModule with a batch dimension. How to type hint?"""
    ....

#
# Runtime code
#
import jax.random as jr

@eqx.filter_vmap
def make_module(a, b):
    return TestModule(a, b)

dim = 10
key = jr.key(1234)
a = jr.random.normal(key, shape=(dim,))
b = jr.random.normal(key, shape=(dim, 2))

vmapped_module  = make_module(a, b)
compute_something(vmapped_module)

How would I type hint the compute_something function, which specifically takes in a TestModule with particular batch dimensions?

Start with the TestModule, compute_something, and eqx.filter_vmap examples in the issue, then review existing jaxtyping guidance for shaped arrays and JAX or Equinox integration. Done means establishing whether a supported annotation can express the vmapped TestModule's batch dimensions and documenting or implementing the agreed approach.