Bug Report #5203
Description
Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.
System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
- Flax, jax, jaxlib versions (obtain with
pip show flax jax jaxlib: - Python version:
- GPU/TPU model and memory:
- CUDA version (if applicable):
Problem you have encountered:
What you expected to happen:
Logs, error messages, etc:
Steps to reproduce:
Whenever possible, please provide a minimal example. Please consider submitting it as a Colab link.