einops-like packing notation

Hey @patrick-kidger,

I'm wondering how hard it would be to have einops-like notation for packed axes? For example,

Float[Array, "B C (H W)"]

would indicate that the last axis is a flattened version of the height and width axis.

This means that if you have the full signature as:

def unpack(x: Float[Array, "B C (H W)"]) -> Float[Array, "B C H W"]:
    ... # magic unpacking
    return y

then jaxtyping would check that y.shape[2] * y.shape[3] == x.shape[2].

Note that in many cases it would not be able to confirm H and W individually. I think that is okay; it's just free variables. But if it can confirm the individual shapes, then it can do the type check.

What do you think? Does this make sense?

The issue names no implementation files or tests. Start by tracing the existing runtime shape-checking path for Python annotations, then determine how grouped axes such as (H W) should be parsed and checked. Done means supporting the example notation and validating the product of grouped dimensions when individual dimensions cannot be inferred.