Comments (4)
A first recourse would be to report this as a bug in Flax. I suspect they'd be amenable to a PR fixing this.
An alternative option is to use Equinox instead, of course!
How are you applying jaxtyping's typechecking? For the more general problem of skipping certain modules then usually the import hook will only add those modules you explicitly ask it to typecheck. (Where you can list your own modules and not list third-party modules.)
from jaxtyping.
An alternative option is to use Equinox instead, of course!
If only it were that easy, migrating the entire codebase to Equinox is a bit out of scope unfortunately 😅
Upon further inspection, it looks like this only happens when typechecking a flax.struct.dataclass
.
I'm using the decorator syntax to only typecheck certain functions and dataclasses, so I'm not sure how the import hook would help. The problem is that jaxtyping is working great most of the time (during init/apply/etc), and I would really like to typecheck my dataclasses at those times! It just breaks during tabulate. Are there any other ways that you know of to disable typechecking just for that one call?
I also tried defining Array = jax.Array | flax.linen.summary._ArrayRepresentation
and using that type annotation instead, to no avail.
from jaxtyping.
Ah sorry, I had to go look up the source code for tabulate
to understand what you're referring to. IIUC they're actually initialising the dataclass for you, but with the arguments of the wrong type?
For what it's worth, jaxtyping (for historical reasons) actually checks the attributes of the class after initialisation, not the arguments to __init__
. When you report things not working above, it might be that you're doing this for the arguments of a custom __init__
despite it being the types of the fields that you'd need to change? (To e.g. Float[Array, "foo bar"] | flax.linen.summary._ArrayRepresentation
.)
Other than that jaxtyping does offer an opt-out flag that you could set for the duration of the call:
jaxtyping/jaxtyping/_config.py
Line 28 in 49cf97e
from jaxtyping.
Ah thanks, the opt-out flag is the escape hatch I'm looking for!
It seems like the reason jax.Array | flax.linen.summary._ArrayRepresentation
didn't work, though, is a separate issue with Unions. Let me open a separate issue.
from jaxtyping.
Related Issues (20)
- jax dependency error when jax is not installed HOT 4
- Random instances / Hypothesis-like generation HOT 3
- Question: manual assertion HOT 4
- Move equinox "tree_pformat" into jaxtyping or allow users to configure their own HOT 1
- How to use with Sphinx autodoc? HOT 1
- Issues with torch.compile HOT 5
- Functions without type hints and import hook HOT 1
- Compatibility with multiprocessing / joblib - AttributeError: type object 'Float[Tensor, 'batch_size num_classes']' has no attribute 'index_variadic' HOT 4
- Can typeguard be an optional dependency? HOT 8
- Are pytorch named tensors supported, like in torchtyping? HOT 1
- How to properly escape `*` and `_` when rendering docs with Sphinx HOT 3
- numpy structured dtype support HOT 1
- Bug with default argument binding HOT 2
- Unions not working HOT 4
- jaxtyping with JAX severely slowing down training speed HOT 6
- `install_import_hook` skip `no_type_check` HOT 1
- add type annotation for jaxtyped HOT 1
- Installing `jaxtyping` makes Pytest slow(er) HOT 3
- Annotations for tensors with dynamics dimensions HOT 2
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from jaxtyping.