flax.linen.module.init fails under dynamic type checking

[SalamanderXing]

Colab Link

The following MRE fails:

from jax import numpy as jnp
import flax.linen as nn
import jax
from beartype import beartype


@beartype
class MyModule(nn.Module):
    @nn.compact
    def __call__(self, x):
        return nn.Dense(1)(x)


model = MyModule()

params = model.init(
    rngs={"params": jax.random.PRNGKey(0)},
    x=jnp.ones((1, 1)),
)

With the following stack trace:

/usr/local/lib/python3.10/dist-packages/beartype/_util/hint/pep/utilpeptest.py:311: BeartypeDecorHintPep585DeprecationWarning: PEP 484 type hint typing.Type[flax.core.scope.Scope] deprecated by PEP 585. This hint is scheduled for removal in the first Python version released after October 5th, 2025. To resolve this, import this hint from "beartype.typing" rather than "typing". For further commentary and alternatives, see also:
    https://beartype.readthedocs.io/en/latest/api_roar/#pep-585-deprecations
  warn(
/usr/local/lib/python3.10/dist-packages/beartype/_util/hint/pep/utilpeptest.py:311: BeartypeDecorHintPep585DeprecationWarning: PEP 484 type hint typing.Type[flax.linen.module.Module] deprecated by PEP 585. This hint is scheduled for removal in the first Python version released after October 5th, 2025. To resolve this, import this hint from "beartype.typing" rather than "typing". For further commentary and alternatives, see also:
    https://beartype.readthedocs.io/en/latest/api_roar/#pep-585-deprecations
  warn(
/usr/local/lib/python3.10/dist-packages/beartype/_util/hint/pep/utilpeptest.py:311: BeartypeDecorHintPep585DeprecationWarning: PEP 484 type hint typing.Type[flax.linen.module._Sentinel] deprecated by PEP 585. This hint is scheduled for removal in the first Python version released after October 5th, 2025. To resolve this, import this hint from "beartype.typing" rather than "typing". For further commentary and alternatives, see also:
    https://beartype.readthedocs.io/en/latest/api_roar/#pep-585-deprecations
  warn(
WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
---------------------------------------------------------------------------
BeartypeCallHintParamViolation            Traceback (most recent call last)
[<ipython-input-3-197af075d233>](https://localhost:8080/#) in <cell line: 16>()
     14 model = MyModule()
     15 
---> 16 params = model.init(
     17     rngs={"params": jax.random.PRNGKey(0)},
     18     x=jnp.ones((1, 1)),

    [... skipping hidden 8 frame]

<@beartype(__main__.MyModule.__init__) at 0x7b553c655bd0> in __init__(__beartype_func, __beartype_conf, __beartype_get_violation, __beartype_object_99245904980800, __beartype_object_99245939285840, __beartype_object_99245939373632, __beartype_object_99245939385792, __beartype_object_135606015798272, *args, **kwargs)

BeartypeCallHintParamViolation: Method __main__.MyModule.__init__() parameter parent=<flax.core.scope.Scope object at 0x7b553c4da2f0> violates type hint typing.Union[typing.Type[flax.linen.module.Module], typing.Type[flax.core.scope.Scope], typing.Type[flax.linen.module._Sentinel], NoneType], as <class "flax.core.scope.Scope"> <flax.core.scope.Scope object at 0x7b553c4da2f0> not type or <class "builtins.NoneType">.

The issue might be that the flax.linen.module.init function is wrongly annotated.

Related Discussion

[chiamp]

Thanks for catching this!