Symbolic expressions now support delayed binding to arguments. Fixes #93.

[patrick-kidger]

No description provided.

[dlwh]

pretty sure this should be {size}?

[patrick-kidger]

Agreed! I'll fix it now.
This does make me think, though -- handling a general shape is a little annoying with this change:

def full(shape: tuple[int, ...], fill: float) -> Float[Array, "{strshape}"]:
    strshape = " ".join(map(str, shape))
    return jnp.full(shape, fill)

I don't have a good answer for that right now.

[dlwh]

yeah i thought about this for a minute and decided it was tough. maybe just a helper function? jaxtyping.strshape that just does the first line? tuple shapes are pretty common!

[danielward27]

Not sure if this is intended, but after reading the updated docs it isn't obvious to me why one of these would work, and the other would not

import equinox as eqx
import jax.numpy as jnp
from beartype import beartype as typechecker
from jaxtyping import ArrayLike, Shaped, jaxtyped


class SomeClass1(eqx.Module):
    dim: int

    @jaxtyped
    @typechecker
    def __call__(self, arr: Shaped[ArrayLike, "2 {self.dim}"]):
        return arr


class SomeClass2(eqx.Module):
    dim: int

    @jaxtyped
    @typechecker
    def __call__(self, arr: Shaped[ArrayLike, "{self._strshape}"]):
        return arr

    @property
    def _strshape(self):
        return f"2 {self.dim}"


dim = 5
x = jnp.ones((2, dim))
my_class1 = SomeClass1(dim)
my_class2 = SomeClass2(dim)


my_class1(x)  # This works
my_class2(x)  # This raises BeartypeCallHintParamViolation

[patrick-kidger]

Ah, that's an interesting edge-case. The reason is that delayed bindings of the form "{self.foo}" can only be used for a single name (whether that's a single axis e.g. batch or variadic axes e.g. *batch). Here you're trying to expand to two names: 2 5.

I suspect that's probably a limitation that can be lifted though! I'll see what I can do.

[patrick-kidger]

Having now looked at this a bit, I don't think this is something that can easily be arranged. Right now the code for separating out each axis runs when the annotation is first created. Meanwhile, the code for these kinds of f-strings is ran during the isinstance check.

So in order to support use-cases like the above, we would need to delay running the "axis handling" code until after the f-string handling code. This would (a) represent a small performance drop for unJIT'd code, and (b) be a fairly fiddly rewrite to arrange.

As such I think this probably won't be something supported, at least in the near future.

[danielward27]

I see, that seems reasonable. Thanks for having a look anyway! I guess this means objects with tuple[int, ...] shape attributes can't use their shapes in the parameter annotations too (unless you know the length of the shape beforehand, in which case you could manually splat it out).