Merge pull request #2446 from jheek:remove-pytype-generic-workaround

PiperOrigin-RevId: 473235778
google · Sep 9, 2022 · e320e11 · e320e11
2 parents 8687673 + 70d845a
commit e320e11
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]):