nested NamedTuples don't work with beartype + jaxtyping

This works

from beartype import beartype as typechecker
from jaxtyping import Array, Float, jaxtyped
from typing import NamedTuple, Union, Optional

@typechecker
class ParamsLGSSMInitial(NamedTuple):
    mean: float 
    cov: float 

p = ParamsLGSSMInitial(mean=5.0, cov=10.0) # fine

@typechecker
class ParamsLGSSM(NamedTuple): # fine
    initial: ParamsLGSSMInitial

but this does not

@jaxtyped
@typechecker
class ParamsLGSSMInitial2(NamedTuple):
    mean: float 
    cov: float 

p = ParamsLGSSMInitial2(mean=5.0, cov=10.0) # fine

@typechecker
class ParamsLGSSM(NamedTuple): 
    initial2: ParamsLGSSMInitial2 # fail

I get the error

beartype.roar.BeartypeDecorHintNonpepException: @beartyped namedtuple_ParamsLGSSM.ParamsLGSSM.__new__() parameter "initial2" type hint <jaxtyping.decorator._Jaxtyped object at 0x7f7ae90450a0> either PEP-noncompliant or currently unsupported by @beartype.

I realize this may be a problem with jaxtyping and not beartype… (Also I realize I don’t need the @jaxtyped decorator in this example, but in the real code, float is actually an Array.)

@patrick-kidger @slinderman

[patrick-kidger]

So for context,jaxtyped(fn) returns a functools.wraps(fn)(_Jaxtyped(fn)). I feel like that should still be a valid type annotation (whether or not beartype respects the wraps).

FWIW one possible fix might be to make jaxtyped of a class actually just wrap the __init__ method.

I’d be happy to accept a PR on this over in the jaxtyping repo.

[leycec]

I think this should be fixed…

A design goal for jaxtyping is to avoid the need for that kind of close coupling!

The sweet words I’ve waited a lifetime to hear. 🥰 😍 😻 :heart_on_fire:

The ML Dream Team of @murphyk + @patrick-kidger capably resolved this without my needing to lift a paw. Thanks so much for resolving this for me. I can now quietly close this and pretend I knew what was going on here.