Possible bug when using `torch.nn.Module` and `@jaxtyping` + `@typechecker` about jaxtyping HOT 5 CLOSED

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.

from jaxtyping.

They both work.

from jaxtyping.

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.

from jaxtyping.

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!

from jaxtyping.

Closing as I think this issue is resolved.

from jaxtyping.