Is it possible to be used with tensordict?

[RuiWang1998]

I find jaxtyping a blast to use, as is tensordict(https://github.com/pytorch/tensordict). If I could have it both ways, that would be even more amazing!

Thanks for the great work!

[patrick-kidger]

Hmm, so it's not clear to me how we'd go about annotating a TensorDict, in the same way that it's not obvious to me how to annotate the shapes and dtypes of the values of a regular dict.

But... have you tried the @tensorclass decorator that tensordict also provides? As that explicitly annotates the contained elements then I think that should be exactly what you need?

[Mxbonn]

Can confirm that tensorclass works

@tensorclass
class CameraIntrinsics:
    fx: Float[Tensor, " *num_cameras"]
    fy: Float[Tensor, " *num_cameras"]
    cx: Float[Tensor, " *num_cameras"]
    cy: Float[Tensor, " *num_cameras"]
...    
def some_other_fun(intrinsics: Shaped[CameraIntrinsics, " *num_cameras"])

[RuiWang1998]

Great! Thanks for the advice!