Dimension changing

[samuelstevens]

import beartype
import torch

from jaxtyping import Float, jaxtyped
from torch import Tensor


@beartype.beartype
class Identity(torch.nn.Module):
    @jaxtyped(typechecker=beartype.beartype)
    def forward(self, x: Float[Tensor, "batch d_model"]):
        return x


@jaxtyped(typechecker=beartype.beartype)
def main():
    model = Identity()

    x = torch.ones((4, 1), dtype=torch.float)
    model(x)

    x = torch.ones((6, 1), dtype=torch.float)
    model(x)  # <- Throws an error about "batch" not matching.


if __name__ == "__main__":
    main()

Not sure why this is happening. But this second call to model.forward will complain because "batch" has changed between calls. How can I fix this?

[patrick-kidger]

I think the additional beartype decorator you have on the class is responsible. If I remove this then the code passes successfuly.

[samuelstevens]

Yeah, and it also passes if I replace it with jaxtyped(typechecker=beartype.beartype). Not sure why this is the case, or how to think about this so I can avoid it in the future. If you have any ideas, that would be appreciated. If not, feel free to close this issue. Thanks!