[Documentation] Comparison with `tensor_annotations` about jaxtyping HOT 9 CLOSED

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?

1. The name originated from a time when only JAX arrays were supported.
2. 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.)
3. Easier to market!

For completeness, the main differences are that:

1. 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 contrast jaxtyping's static checking only handles the "is this a tensor or not?" part.

2. 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 like ArrayLike, ShapeLike, DtypeLike etc. Also some JAX-specific stuff like Tracers. (=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 of jax.typing, since right now the things that jaxtyping 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.
• Annotate NumPy tensors with either:
  • For typing shapes, the official numpy.typing API.
  • For typing shapes and additional metadata, the third-party nptyping API.
• Annotate PyTorch tensors with the third-party TorchTyping API.
• Annotate TensorFlow tensors with the third-party tensor_annotations API subject to the extreme caveat that tensor_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.