Perform addition/multiplication on shape information about jaxtyping HOT 1 CLOSED

Haha yup, this has always been on the wish-list of those doing dynamic array shape-checking.

One of the complexities here is what to do if someone writes e.g.

def f(x: f["2*dim"]) -> f["3*dim"]): ...

or worse

def f(x: f["dim+(dim**2)"]): ...

where in general we'd have to solve an arbitrary algebraic expression simply to figure out what dim is, and whether all of the arguments are consistent.

Nonetheless, as of version 0.1.0, a version of this feature is now supported!

We take a pragmatic approach, and simply mandate that all dimensions must have had a "normal" name binding by the time they occur in a symbolic expression. Typically this means that you have to use "normal" name bindings in the annotations for the arguments, but are then free to use these in arbitrary symbolic expressions in the annotation for the return type. In practice I imagine that's 99% of all use-cases for this feature any way.

from jaxtyping.