Remove pytype generic workaround

google · Sep 7, 2022 · 70d845a · 70d845a
1 parent fdd1d6f
commit 70d845a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flax/core/lift.py b/flax/core/lift.py
@@ -315,13 +315,13 @@ def swap(target):
 @dataclasses.dataclass(frozen=True)
 class In(Generic[T]):
   """Specifies a variable collection should only be lifted as input."""
-  axis: Any  # pytype does not support generic variable annotation
+  axis: T
 
 
 @dataclasses.dataclass(frozen=True)
 class Out(Generic[T]):
   """Specifies a variable collection should only be lifted as output."""
-  axis: Any  # pytype does not support generic variable annotation
+  axis: T
 
 
 def _split_in_out_axes(xs: Mapping[CollectionFilter, Any]):