Comments (9)
You graciously name-dropped @beartype in your README.md
a few months ago. Now, @beartype humbly returns the favour with:
- A new FAQ entry on JAX typing, which advises everyone to just use
jaxtyping
already. - A new FAQ entry on PyTorch typing, which advises everyone to just use TorchTyping already.
Of course, @beartype documentation is still an unreadable monolithic mess of reStructuredText (reST). Therefore, this means less than it should. But thanks so much for all the phenomenal work you do here, at TorchTyping, and probably elsewhere! Google got the best with @patrick-kidger. 👏 👏
from jaxtyping.
Closing this issue! Current state of affairs:
jaxtyping 🤝 beartype
from jaxtyping.
Haha, thank you :D
So actually, despite the name -- jaxtyping
also covers TensorFlow, NumPy, and PyTorch! Not just JAX. Mostly because I could and it was easy, but also because it's actually pretty common to use TF/NP/PT dataloaders prior to a JAX computation.
So clearly jaxtyping
's support here should be better-advertised!
Given this, why the jaxtyping
name?
- The name originated from a time when only JAX arrays were supported.
- I'm not willing to commit to supporting every tensor/array framework under the sun. (The main discrepancy is how they all handles dtypes differently.)
- Easier to market!
For completeness, the main differences are that:
-
tensor_annotations
went down the hellishly complicated/unmaintainable approach of actually adding type annotations to the TF and JAX APIs, and this means that they're able to do shape-checking statically. In contrastjaxtyping
's static checking only handles the "is this a tensor or not?" part. -
jaxtyping
instead provides support for checking many more things (at runtime only): variadic dimensions, broadcasting, pytrees, etc. (Also it's actually maintainable!)
from jaxtyping.
Well, "always use jaxtyping" certainly sounds like an excellent life choice to me. :)
To suppory this use case, I should probably make the JAX dependence optional. Which in practice should be pretty easy; just lazy-import it for Array
and PyTree
, which are the only pieces with an explicit dependence on JAX.
Regarding the naming -- this is actually quite a careful choice. The intended scope of this project can be summarised in the following diagram:
| array types | other types
--------+-------------+-------------
JAX | yes | yes
not JAX | no | no
as jaxtyping already also provides support for JAX-specific notions like PyTrees, and may in the future also feature JAX-specific notions of DtypeLike
etc.
I actually intend to completely overhaul TorchTyping at some point, and make it a fork of the relevant parts of jaxtyping. See also #18.
I don't think a multi-backend frankenpackage is a good fit for this problem. For example, jaxtyping
currently provides no support for declaring what device an array is on, because that's not a notion that exists in JAX. TorchTyping does, because this is something that exists in PyTorch.
from jaxtyping.
jax.typing
is only for typing array-ness, and also things likeArrayLike
,ShapeLike
,DtypeLike
etc. Also some JAX-specific stuff likeTracer
s. (=dummy arrays threaded through to compile a program.) In particular typing shapes and dtypes is out-of-scope.jaxtyping
is intended to be a superset ofjax.typing
, since right now the things thatjaxtyping
does are still considered too magic for JAX core.
But otherwise, yup, you've pretty much got it!
from jaxtyping.
Sweet tensor Jesus. So what you're saying is everyone should use jaxtyping
and no one should use anything else. Well, alrighty then. I'll go ahead and emphasize that in @beartype's documentation as well when I find a spare weekend lying around.
tensor_annotations
went down the hellishly complicated/unmaintainable approach of actually adding type annotations to the TF and JAX APIs...
Noticed that. Then I shuddered. That doesn't scale to an arbitrary number of competing tensor frameworks... unlike jaxtyping
! Yet again, less
is more
.
Apparently, I'm supposed to be using most
instead of less
now as my default pager. Ohhhhh, no. They can pry less
from my cold Canadian fingers wracked with arthritis.
jaxtyping
instead provides support for checking many more things (at runtime only)
Runtime. This is the way.
Also it's actually maintainable!
I... I like what I'm hearing. Your words are convincing. Your code, even more convincing. Clearly, at least one of us is in a right frame of mind.
Given this, why the
jaxtyping
name?
Gah! This is the only contentious point I contest. It's the curmudgeon in me that's makings me say this, but... it might be useful to generalize the name at some point.
I know. I know. jaxtyping
already has mind share and is a trending thing. But this isn't just JAX typing anymore, is it? Right? Your codebase grew up before anyone noticed. It's like kids and cats that way.
Technically, you'd rather not commit to supporting every existing tensor framework. Yes. I wouldn't either. There's meaningful work to do – and that's not it. Pragmatically, you've already committed to supporting JAX, TensorFlow, NumPy, and PyTorch. So, basically every existing tensor framework. Nobody talks about those other guys.
jaxtyping
→ radtensorhints
is what I'm suggesting. 🧐
from jaxtyping.
Today I learned... that the more I know, the less I know. Because the typing landscape is moving so rapidly, this list summarizes the little that I think I know:
- Annotate JAX tensors with either:
- For typing shapes,...and possibly other things? the official
jax.typing
API – which technically doesn't exist yet, except by the time I eventually document this it probably will. - For typing shapes and additional metadata, the third-party
jaxtyping
API.
- For typing shapes,...and possibly other things? the official
- Annotate NumPy tensors with either:
- For typing shapes, the official
numpy.typing
API. - For typing shapes and additional metadata, the third-party
nptyping
API.
- For typing shapes, the official
- Annotate PyTorch tensors with the third-party TorchTyping API.
- Annotate TensorFlow tensors with the third-party
tensor_annotations
API subject to the extreme caveat thattensor_annotations
does not support runtime type-checking. Does Google really have no off-the-shelf solution for typing TensorFlow in Python? If not, I cry.
Does that fully capture best typing practices for tensors as of Halloween? If so, mash that pumpkin. @beartype will freely shower the above advice onto data scientists in our FAQ. 🎃
I don't think a multi-backend frankenpackage is a good fit for this problem.
But I still want a multi-backend frankenpackage! Cue a new beartype.vale.IsTensor[...]
API. It's only a matter of time. Literally. I don't have the time, so it's probably no matter.
from jaxtyping.
In particular typing shapes and dtypes is out-of-scope.
Gotchas. I badly assumed jax.typing
was an orthogonal copycat of numpy.typing
and thus typed dtypes like numpy.typing.NDArray[{dtype}]
. Assumptions will be the death of @beartype docos.
Thanks so much for the clarifications! I'll get this typed up when life throws me a bucket of free time. 😓
from jaxtyping.
Hurrah! Thank you - that's really cool to see. 😎
from jaxtyping.
Related Issues (20)
- bug: can't type flax.struct.dataclass with vmapped functions HOT 3
- Disabling JAX import HOT 2
- einops-like packing notation HOT 4
- Symbolic expressions example doesn't run HOT 2
- v0.2.26 Release Failed HOT 1
- Allow two variadic shapes when it makes sense HOT 1
- Support for NestedTensors HOT 3
- `jaxtyped` Annotation fails
- 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
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.