Remove typeguard checks

pyproject.toml

@@ -49,6 +49,7 @@ ignore = [
     "E501",  # Line too long.
     "E721",  # Do not compare types, use `isinstance()`.
     "F722",  # Forward annotation false positive from jaxtyping. Should be caught by pyright.
+    "F731",  # Do not assign a lambda expression, use a def.
     "F821",  # Forward annotation false positive from jaxtyping. Should be caught by pyright.
     "PLR2004",  # Magic value used in comparison.
     "PLR0915",  # Too many statements.

src/egoallo/data/dataclass.py

@@ -5,16 +5,14 @@
 import numpy as np
 import torch
 import torch.utils.data
-import typeguard
-from jaxtyping import Bool, Float, jaxtyped
+from jaxtyping import Bool, Float
 from torch import Tensor
 
 from .. import fncsmpl, fncsmpl_extensions
 from .. import transforms as tf
 from ..tensor_dataclass import TensorDataclass
 
 
-@jaxtyped(typechecker=typeguard.typechecked)
 class EgoTrainingData(TensorDataclass):
     """Dictionary of tensors we use for EgoAllo training."""
 

src/egoallo/fncsmpl.py

@@ -26,16 +26,14 @@
 
 import numpy as np
 import torch
-import typeguard
 from einops import einsum
-from jaxtyping import Float, Int, jaxtyped
+from jaxtyping import Float, Int
 from torch import Tensor
 
 from .tensor_dataclass import TensorDataclass
 from .transforms import SE3, SO3
 
 
-@jaxtyped(typechecker=typeguard.typechecked)
 class SmplhModel(TensorDataclass):
     """A human body model from the SMPL family."""
 
@@ -118,7 +116,6 @@ def with_shape(self, betas: Float[Tensor, "*#batch n_betas"]) -> SmplhShaped:
         )
 
 
-@jaxtyped(typechecker=typeguard.typechecked)
 class SmplhShaped(TensorDataclass):
     """The SMPL-H body model with a body shape applied."""
 
@@ -183,7 +180,6 @@ def with_pose(
         )
 
 
-@jaxtyped(typechecker=typeguard.typechecked)
 class SmplhShapedAndPosed(TensorDataclass):
     shaped_model: SmplhShaped
     """Underlying shaped body model."""
@@ -266,7 +262,6 @@ def lbs(self) -> SmplMesh:
         )
 
 
-@jaxtyped(typechecker=typeguard.typechecked)
 class SmplMesh(TensorDataclass):
     """Outputs from the SMPL-H model."""