We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
After update from 0.2.0: AttributeError: module 'jax.core' has no attribute 'eval_context'
After updating from flax 0.2.0 to flax 0.2.2 I get the above error message. Downgrading to 0.2.0 solves this, so the error source is located. I'm working with the now deprecated flax.nn package if backward-compatibility might be the reason for this issue.
The Issue is encountered in a custom RNN, when using the init_by_shape function in conjunction with jax.lax.scan.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `setup.py`
Content:
```
1 # Copyright 2020 The Flax Authors.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """setup.py for Flax."""
16 
17 import os
18 from setuptools import find_packages
19 from setuptools import setup
20 
21 here = os.path.abspath(os.path.dirname(__file__))
22 try:
23   README = open(os.path.join(here, "README.md"), encoding='utf-8').read()
24 except IOError:
25   README = ""
26 
27 install_requires = [
28     "numpy>=1.12",
29     "jax>=0.1.59",
30     "matplotlib",  # only needed for tensorboard export
31     "dataclasses;python_version<'3.7'", # will only install on py3.6
32     "msgpack",
33 ]
34 
35 tests_require = [
36     "atari-py",
37     "clu",  # All examples.
38     "gym",
39     "jaxlib",
40     "ml-collections",
41     "opencv-python",
42     "pytest",
43     "pytest-cov",
44     "pytest-xdist==1.34.0",  # upgrading to 2.0 broke tests, need to investigate
45     "sentencepiece",  # WMT example.
46     "svn",
47     "tensorflow",
48     "tensorflow_text",  # WMT example.
49     "tensorflow_datasets",
50 ]
51 
52 __version__ = None
53 
54 with open('flax/version.py') as f:
55   exec(f.read(), globals())
56 
57 setup(
58     name="flax",
59     version=__version__,
60     description="Flax: A neural network library for JAX designed for flexibility",
61     long_description="\n\n".join([README]),
62     long_description_content_type='text/markdown',
63     classifiers=[
64         "Development Status :: 3 - Alpha",
65         "Intended Audience :: Developers",
66         "Intended Audience :: Science/Research",
67         "License :: OSI Approved :: Apache Software License",
68         "Programming Language :: Python :: 3.7",
69         "Topic :: Scientific/Engineering :: Artificial Intelligence",
70         ],
71     keywords="",
72     author="Flax team",
73     author_email="flax-dev@google.com",
74     url="https://github.com/google/flax",
75     packages=find_packages(),
76     include_package_data=False,
77     zip_safe=False,
78     install_requires=install_requires,
79     extras_require={
80         "testing": tests_require,
81         },
82     )
83
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -26,7 +26,7 @@
 
 install_requires = [
     "numpy>=1.12",
-    "jax>=0.1.59",
+    "jax>=0.1.77",
     "matplotlib",  # only needed for tensorboard export
     "dataclasses;python_version<'3.7'", # will only install on py3.6
     "msgpack",

{"golden_diff": "diff --git a/setup.py b/setup.py\n--- a/setup.py\n+++ b/setup.py\n@@ -26,7 +26,7 @@\n \n install_requires = [\n     \"numpy>=1.12\",\n-    \"jax>=0.1.59\",\n+    \"jax>=0.1.77\",\n     \"matplotlib\",  # only needed for tensorboard export\n     \"dataclasses;python_version<'3.7'\", # will only install on py3.6\n     \"msgpack\",\n", "issue": "After update from 0.2.0: AttributeError: module 'jax.core' has no attribute 'eval_context'\nAfter updating from flax 0.2.0 to flax 0.2.2 I get the above error message. Downgrading to 0.2.0 solves this, so the error source is located. I'm working with the now deprecated flax.nn package if backward-compatibility might be the reason for this issue.\r\nThe Issue is encountered in a custom RNN, when using the init_by_shape function in conjunction with jax.lax.scan.\n", "before_files": [{"content": "# Copyright 2020 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"setup.py for Flax.\"\"\"\n\nimport os\nfrom setuptools import find_packages\nfrom setuptools import setup\n\nhere = os.path.abspath(os.path.dirname(__file__))\ntry:\n  README = open(os.path.join(here, \"README.md\"), encoding='utf-8').read()\nexcept IOError:\n  README = \"\"\n\ninstall_requires = [\n    \"numpy>=1.12\",\n    \"jax>=0.1.59\",\n    \"matplotlib\",  # only needed for tensorboard export\n    \"dataclasses;python_version<'3.7'\", # will only install on py3.6\n    \"msgpack\",\n]\n\ntests_require = [\n    \"atari-py\",\n    \"clu\",  # All examples.\n    \"gym\",\n    \"jaxlib\",\n    \"ml-collections\",\n    \"opencv-python\",\n    \"pytest\",\n    \"pytest-cov\",\n    \"pytest-xdist==1.34.0\",  # upgrading to 2.0 broke tests, need to investigate\n    \"sentencepiece\",  # WMT example.\n    \"svn\",\n    \"tensorflow\",\n    \"tensorflow_text\",  # WMT example.\n    \"tensorflow_datasets\",\n]\n\n__version__ = None\n\nwith open('flax/version.py') as f:\n  exec(f.read(), globals())\n\nsetup(\n    name=\"flax\",\n    version=__version__,\n    description=\"Flax: A neural network library for JAX designed for flexibility\",\n    long_description=\"\\n\\n\".join([README]),\n    long_description_content_type='text/markdown',\n    classifiers=[\n        \"Development Status :: 3 - Alpha\",\n        \"Intended Audience :: Developers\",\n        \"Intended Audience :: Science/Research\",\n        \"License :: OSI Approved :: Apache Software License\",\n        \"Programming Language :: Python :: 3.7\",\n        \"Topic :: Scientific/Engineering :: Artificial Intelligence\",\n        ],\n    keywords=\"\",\n    author=\"Flax team\",\n    author_email=\"flax-dev@google.com\",\n    url=\"https://github.com/google/flax\",\n    packages=find_packages(),\n    include_package_data=False,\n    zip_safe=False,\n    install_requires=install_requires,\n    extras_require={\n        \"testing\": tests_require,\n        },\n    )\n", "path": "setup.py"}], "after_files": [{"content": "# Copyright 2020 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"setup.py for Flax.\"\"\"\n\nimport os\nfrom setuptools import find_packages\nfrom setuptools import setup\n\nhere = os.path.abspath(os.path.dirname(__file__))\ntry:\n  README = open(os.path.join(here, \"README.md\"), encoding='utf-8').read()\nexcept IOError:\n  README = \"\"\n\ninstall_requires = [\n    \"numpy>=1.12\",\n    \"jax>=0.1.77\",\n    \"matplotlib\",  # only needed for tensorboard export\n    \"dataclasses;python_version<'3.7'\", # will only install on py3.6\n    \"msgpack\",\n]\n\ntests_require = [\n    \"atari-py\",\n    \"clu\",  # All examples.\n    \"gym\",\n    \"jaxlib\",\n    \"ml-collections\",\n    \"opencv-python\",\n    \"pytest\",\n    \"pytest-cov\",\n    \"pytest-xdist==1.34.0\",  # upgrading to 2.0 broke tests, need to investigate\n    \"sentencepiece\",  # WMT example.\n    \"svn\",\n    \"tensorflow\",\n    \"tensorflow_text\",  # WMT example.\n    \"tensorflow_datasets\",\n]\n\n__version__ = None\n\nwith open('flax/version.py') as f:\n  exec(f.read(), globals())\n\nsetup(\n    name=\"flax\",\n    version=__version__,\n    description=\"Flax: A neural network library for JAX designed for flexibility\",\n    long_description=\"\\n\\n\".join([README]),\n    long_description_content_type='text/markdown',\n    classifiers=[\n        \"Development Status :: 3 - Alpha\",\n        \"Intended Audience :: Developers\",\n        \"Intended Audience :: Science/Research\",\n        \"License :: OSI Approved :: Apache Software License\",\n        \"Programming Language :: Python :: 3.7\",\n        \"Topic :: Scientific/Engineering :: Artificial Intelligence\",\n        ],\n    keywords=\"\",\n    author=\"Flax team\",\n    author_email=\"flax-dev@google.com\",\n    url=\"https://github.com/google/flax\",\n    packages=find_packages(),\n    include_package_data=False,\n    zip_safe=False,\n    install_requires=install_requires,\n    extras_require={\n        \"testing\": tests_require,\n        },\n    )\n", "path": "setup.py"}]}