Add **init_kwargs to self.param type annotation.

PiperOrigin-RevId: 622872682
google · Apr 8, 2024 · a1ffa00 · a1ffa00
1 parent 39ce98f
commit a1ffa00
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flax/linen/module.py b/flax/linen/module.py
@@ -1761,7 +1761,7 @@ def variable(
 
  @overload
  def param(
- self, name: str, init_fn: Callable[..., T], *init_args,
+  self, name: str, init_fn: Callable[..., T], *init_args, **init_kwargs
  ) -> T:
  ...