add typing overloads for variable

google · Aug 1, 2023 · 8525a9b · 8525a9b
1 parent 1e2d1c7
commit 8525a9b
Show file tree

Hide file tree

Showing 2 changed files with 89 additions and 3 deletions.
diff --git a/flax/core/scope.py b/flax/core/scope.py
@@ -818,14 +818,57 @@ def put(target, key, val):
 
  put(variables, name, value)
 
+ @overload
+ def variable(
+ self,
+ col: str,
+ name: str,
+ init_fn: Optional[Callable[..., T]] = None,
+ *init_args,
+ ) -> Variable[T]:
+ ...
+
+ @overload
+ def variable(
+ self,
+ col: str,
+ name: str,
+ init_fn: Optional[Callable[..., T]] = None,
+ *init_args,
+ unbox: Literal[True],
+ ) -> Variable[T]:
+ ...
+
+ @overload
+ def variable(
+ self,
+ col: str,
+ name: str,
+ init_fn: Optional[Callable[..., T]] = None,
+ *init_args,
+ unbox: Literal[False],
+ ) -> Variable[meta.AxisMetadata[T]]:
+ ...
+
+ @overload
+ def variable(
+ self,
+ col: str,
+ name: str,
+ init_fn: Optional[Callable[..., T]] = None,
+ *init_args,
+ unbox: bool = True,
+ ) -> Union[Variable[T], Variable[meta.AxisMetadata[T]]]:
+ ...
+
  def variable(
  self,
  col: str,
  name: str, # pylint: disable=keyword-arg-before-vararg
  init_fn: Optional[Callable[..., T]] = None,
  *init_args,
  unbox: bool = True,
- ) -> Variable[T]:
+ ) -> Union[Variable[T], Variable[meta.AxisMetadata[T]]]:
  """Creates a variable if it doesn't exist yet in this scope and returns it.
 
  Args:

diff --git a/flax/linen/module.py b/flax/linen/module.py
@@ -1330,14 +1330,57 @@ def clone_fn(m: Module) -> Module:
 
  return module
 
+ @overload
+ def variable(
+ self,
+ col: str,
+ name: str,
+ init_fn: Optional[Callable[..., T]] = None,
+ *init_args,
+ ) -> Variable[T]:
+ ...
+
+ @overload
+ def variable(
+ self,
+ col: str,
+ name: str,
+ init_fn: Optional[Callable[..., T]] = None,
+ *init_args,
+ unbox: Literal[True],
+ ) -> Variable[T]:
+ ...
+
+ @overload
+ def variable(
+ self,
+ col: str,
+ name: str,
+ init_fn: Optional[Callable[..., T]] = None,
+ *init_args,
+ unbox: Literal[False],
+ ) -> Variable[meta.AxisMetadata[T]]:
+ ...
+
+ @overload
+ def variable(
+ self,
+ col: str,
+ name: str,
+ init_fn: Optional[Callable[..., T]] = None,
+ *init_args,
+ unbox: bool = True,
+ ) -> Union[Variable[T], Variable[meta.AxisMetadata[T]]]:
+ ...
+
  def variable(
  self,
  col: str,
  name: str,
- init_fn: Optional[Callable[..., Any]] = None,
+ init_fn: Optional[Callable[..., T]] = None,
  *init_args,
  unbox: bool = True,
- ) -> Variable:
+ ) -> Union[Variable[T], Variable[meta.AxisMetadata[T]]]:
  """Declares and returns a variable in this Module.
 
  See :mod:`flax.core.variables` for more information. See also :meth:`param`