patrick-kidger / patrick-kidger/jaxtyping
Problem with symbolic shapes
Open
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.9k
- Forks
- 96
- PR merge metrics
- No merged PRs in 30d
Description
This
import jax
from jax import export
from jax import numpy as jnp
import jaxtyping as jt
import typeguard
@jt.jaxtyped(typechecker=typeguard.typechecked)
def f(
x: jt.Float[jt.Array, "*#B"],
) -> jt.Float[jt.Array, "*#B"]:
return x * jnp.sum(x) ** 2
dtype = jnp.float32
x_shape = export.symbolic_shape("b")
export.export(jax.jit(f))(
jax.ShapeDtypeStruct(x_shape, dtype)
)
fails with
jaxtyping.TypeCheckError: Type-check error whilst checking the return value of __main__.f.
Actual value: f32[b](jax)
Expected type: Float[Array, '*#B'].
----------------------
Called with parameters: {'x': f32[b](jax)}
Parameter annotations: (x: Float[Array, '*#B']) -> Any.
The current values for each jaxtyping axis annotation are as follows.
B=(b,)
The problem seems to be *# (it works without either).
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the provided reproducer using export.symbolic_shape, jax.jit, and the @jaxtyped function, then compare the annotation behavior with and without *#. Trace the runtime return-value check and add a regression test showing that symbolic-shape arrays annotated with *# pass type checking.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- tooling
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100