Possible bug when using torch.nn.Module and @jaxtyping + @typechecker

[PaulScemama]

I could be missing something, but I think there may be something problematic happening when decorating functions with @jaxtyping and @typechecker.

Minimal code example showing that when using the decorators @jaxtyping and @typechecker where typechecker is beartype, forward() is no longer a types.MethodType. This causes bugs when trying to train a torch.compile(model) because it needs to assert model.forward is of type types.MethodType.

import torch
import types
import torch.nn as nn
from jaxtyping import jaxtyped, Float
from beartype import beartype as typechecker

class model_without_jaxtyping(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x
    
model_without_jaxtyping_instance = model_without_jaxtyping()
print(type(model_without_jaxtyping_instance.forward) == types.MethodType)
# 'True'

class model_with_jaxtyping(nn.Module):
    def __init__(self):
        super().__init__()

    @jaxtyped
    @typechecker
    def forward(self, x: Float[torch.Tensor, "..."]):
        return x
    
model_with_jaxtyping_instance = model_with_jaxtyping()
print(type(model_with_jaxtyping_instance.forward) == types.MethodType)
# 'False'

[patrick-kidger]

So jaxtyped actually happens to return a class instance, rather than a function. That's the immediate reason for what we're observing here.

For your downstream use-case with torch.compile, can you check whether:

• Whether

  @typechecker
  def forward(...)

  works?

• And whether

  def wrapper(fn):
      @functools.wraps(fn)
      def wrapped(*args, **kwargs):
          return jaxtyped(typechecker(fn))(*args, **kwargs)
      return wrapped
  
  class Model(nn.Module):
      @wrapper
      def forward(...): ...

  works?

That'll give some insight on how we might fix this.

[PaulScemama]

They both work.

[patrick-kidger]

Great. Can you try install from the no-more-jaxtyped branch (c.f. #72) and see if this works for your downstream use-case? If so then I'll merge that PR.

[st--]

I had just come across what is probably the same issue in different clothes when trying to subclass tensorflow_probability's Distribution type - that also got resolved by #72!

[patrick-kidger]

Closing as I think this issue is resolved.