The code below fails, whether I use typeguard or beartype, and gives the error

Yes, that works! <div class="snippet-clipboard-content notranslate position-relati

jaxtyping does not play nicely with inheritance about jaxtyping HOT 2 CLOSED

google commented on May 3, 2024

jaxtyping does not play nicely with inheritance

from jaxtyping.

Comments (2)

patrick-kidger commented on May 3, 2024

Right, this is because jaxtyped returns a new object, different to the underlying type.

(It was only designed to work on functions.)

We could probably extend it to classes in the following way:

def jaxtyped(fn):
    if inspect.isclass(fn):
        init = jaxtyped(fn.__init__)
        fn.__init__ = init
        return fn
    else:
        ... # existing implementation

Give that a try and see if it works? If it does then I would be happy to accept the above as a PR.

from jaxtyping.

murphyk commented on May 3, 2024

Yes, that works!


from jaxtyping import  jaxtyped
#from beartype import beartype as typechecker
from typeguard import typechecked as typechecker
import inspect

def jaxtyped2(fn):
    if inspect.isclass(fn):
        init = jaxtyped(fn.__init__)
        fn.__init__ = init
        return fn
    else:
        jaxtyped(fn) # existing implementation
       

@jaxtyped2
@typechecker
class LinearGaussianSSM():
    def __init__(
        self,
        state_dim: int,
        emission_dim: int,
        input_dim: int=0,
        has_dynamics_bias: bool=True,
        has_emissions_bias: bool=True
    ):
        self.state_dim = state_dim
        self.emission_dim = emission_dim
        self.input_dim = input_dim
        self.has_dynamics_bias = has_dynamics_bias
        self.has_emissions_bias = has_emissions_bias



    
@typechecker
class LinearGaussianConjugateSSM(LinearGaussianSSM):
    def __init__(self,
                 state_dim,
                 emission_dim,
                 input_dim=0,
                 has_dynamics_bias=True,
                 has_emissions_bias=True):
        super().__init__(state_dim=state_dim, emission_dim=emission_dim, input_dim=input_dim,
        has_dynamics_bias=has_dynamics_bias, has_emissions_bias=has_emissions_bias)

from jaxtyping.

Recommend Projects