beartype around some __init__

pi_zero_pytorch/pi_zero.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
-from typing import Callable
+
+from beartype.typing import Callable
+from beartype import beartype
+
 from functools import partial
 
 import torch
@@ -20,6 +23,7 @@
 
 from pi_zero_pytorch.tensor_typing import Float, Int, Bool
 
+
 import tqdm
 
 # ein notation
@@ -95,6 +99,7 @@ def direction_loss(pred, target, dim = -1):
 # attention
 
 class Attention(Module):
+    @beartype
     def __init__(
         self,
         dim,
@@ -332,6 +337,7 @@ def forward(self, actions, cond):
 # main class
 
 class PiZero(Module):
+    @beartype
     def __init__(
         self,
         dim,

pyproject.toml

@@ -25,6 +25,7 @@ classifiers=[
 ]
 
 dependencies = [
+    "beartype",
     "einx>=0.3.0",
     "einops>=0.8.0",
     "jaxtyping",