[Question] Integrate something like jaxtyping for more specific type annotations and runtime type-checking

[dantp-ai]

Question

Are there any plans on the roadmap to use something like jaxtyping for more specific type annotations and runtime type checking?

Additional context

Checklist

• I have read the documentation (required)
• I have checked that there is no similar issue in the repo (required)

[araffin]

Hello,
there is currently no plan for it.
Would you volunteer for such feature?

[dantp-ai]

I can look into it. I'll start first with the tests folder.

[sandeshkatakam]

Hey,
I was actually considering opening an issue to explore adding jaxtyping support as well. It sounds like we're aligned in our interests, and If you like, I’d love to collaborate with you on this!

I can look into it. I'll start first with the tests folder.

[dantp-ai]

@sandeshkatakam I'd be glad to collaborate. How should we do it best?

One idea would be that, I open a feature branch on my fork. You'll need to add my fork as a remote and check out the same branch. Once this is done, we can sync frequently and decided who works on specific parts to avoid conflicts. What do you think?

[sandeshkatakam]

@sandeshkatakam I'd be glad to collaborate. How should we do it best?

One idea would be that, I open a feature branch on my fork. You'll need to add my fork as a remote and check out the same branch. Once this is done, we can sync frequently and decided who works on specific parts to avoid conflicts. What do you think?

Yeah this seems great! I will add your fork as remote and push commits to the branch, after we decide on dividing the specific parts of the work. Thank you :)

[araffin]

before you do too much work, please open a PR as soon as you have a proof of concept, so we can agree on the details

[sandeshkatakam]

@sandeshkatakam I'd be glad to collaborate. How should we do it best?

One idea would be that, I open a feature branch on my fork. You'll need to add my fork as a remote and check out the same branch. Once this is done, we can sync frequently and decided who works on specific parts to avoid conflicts. What do you think?

Let me know once you fork the repo and create a branch

[dantp-ai]

@sandeshkatakam I'll open the PR at the latest on Thu.

[sandeshkatakam]

@sandeshkatakam I'll open the PR at the latest on Thu.

hey did you open the PR??

[dantp-ai]

@sandeshkatakam I looked in a bit more detail how to use jaxtyping in the project. It can be paired with mypy for static analysis and additionally using e.g. beartype for runtime checks in critical sections (and avoid these runtime checks in performance-critical paths).

I'm blocked this week with some other stuff. Feel free to open a PR (from your fork) if you want to dive in already.