Consider moving away from `torchtyping`
In #1170, the wrong background_color type is something that's hard-ish to notice as a human but should've been easily caught by an IDE; it wasn't flagged because RGBRenderer's constructor uses a torchtyping.TensorType type.
It'd be an intrusive change, but could be worth swapping torchtyping for jaxtyping from the same author, which is typing-checking + IDE-friendly and (despite the name):
- has no JAX dependency as of https://github.com/google/jaxtyping/pull/50
- supports PyTorch tensors
See also https://github.com/google/jaxtyping/issues/18.
Interesting, seems worth the transition. Looks like we could probably be done in chunks, no need to do it all at once. Ill look more into this in a few weeks, but will happily take PRs in the meantime.
Sure, if nobody gets to it we can also chat about it in-person in January.
My feeling is it might be worth automating this change? Seems like almost just a find-and-replace, just with some syntax manipulation + maybe heuristics for determining when to use jaxtyping.Floatvs jaxtyping.Int.
If there’s a script for the transition it could also be run on sdfstudio.