Added support for treepath-dependent sizes.

2026-09-09 11:24:55 +08:00 · 2023-11-27 09:50:02 -08:00
parent 9e1ba8a77d
commit d12291de7e
7 changed files with 296 additions and 64 deletions
@@ -24,6 +24,7 @@ In addition some modifiers can be applied:
    `def add(x: Float[Array, "#foo"], y: Float[Array, "#foo"]) -> Float[Array, "#foo"]`.
 - Prepend `_` to a dimension to disable any runtime checking of that dimension (so that it can be used just as documentation). This can also be used as just `_` on its own: e.g. `"b c _ _"`.
 - Documentation-only names (i.e. they're ignored by jaxtyping) can be handled by prepending a name followed by `=` e.g. `Float[Array, "rows=4 cols=3"]`.
+- Prepend `?` to a dimension to indicate that its size can vary within a PyTree structure. (See [PyTree annotations](../pytree.md).)

 When using multiple modifiers, their order does not matter.