Cristian Garcia

@kaixih it looks good but we will have to wait for JAX to push a new release to pypi for tests to pass (according to your comment).

@kaixih thanks! Seems that `pytest` is still broken.

Based on the output of the test run it seems not: ``` ====== test config ======= PYTEST_OPTS: --cov=flax --cov-report=xml --cov-report=term --cov-config=pyproject.toml RUN_DOCTEST: false RUN_PYTEST: true RUN_MYPY: false RUN_PYTYPE: false GH_VENV:...

As a quick fix maybe add ```yaml venv/bin/python3 -m pip install -U jax jaxlib ``` after this line: https://github.com/google/flax/blob/daf06eadab6c9e9bfb64b30d0623245179965155/.github/workflows/build.yml#L120

Created a PR so you can rebase when merged.

Thanks @canyon289 !

cc @IvyZX

Hey @wenscarl, sorry it took so long. I approved it but you have some errors.

[RNN.__calll__](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.RNN.html#flax.linen.RNN.__call__) has the `seq_lengths` to specify padding elements. We currently don't support masking intermediate elements for performance reasons.

Looking into this...