prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <issue>
  Jax version upgrade (AttributeError: CallPrimitive)
  Using the current version of master 66f9c69 of Haiku, I am getting the following error on Colab
  ```
  AttributeError                            Traceback (most recent call last)
  <ipython-input-3-3a9e6adbfff5> in <module>()
  ----> 1 import haiku as hk
  
  /usr/local/lib/python3.6/dist-packages/haiku/__init__.py in <module>()
       17 
       18 from haiku import data_structures
  ---> 19 from haiku import experimental
       20 from haiku import initializers
       21 from haiku import nets
  
  /usr/local/lib/python3.6/dist-packages/haiku/experimental.py in <module>()
       22 from haiku._src.base import custom_getter
       23 from haiku._src.base import ParamContext
  ---> 24 from haiku._src.dot import to_dot
       25 from haiku._src.lift import lift
       26 from haiku._src.module import profiler_name_scopes
  
  /usr/local/lib/python3.6/dist-packages/haiku/_src/dot.py in <module>()
       23 
       24 from haiku._src import data_structures
  ---> 25 from haiku._src import module
       26 from haiku._src import utils
       27 import jax
  
  /usr/local/lib/python3.6/dist-packages/haiku/_src/module.py in <module>()
       26 from haiku._src import base
       27 from haiku._src import data_structures
  ---> 28 from haiku._src import named_call
       29 from haiku._src import utils
       30 import jax.numpy as jnp
  
  /usr/local/lib/python3.6/dist-packages/haiku/_src/named_call.py in <module>()
       29 
       30 # Registering named call as a primitive
  ---> 31 named_call_p = core.CallPrimitive('named_call')
       32 # named_call is implemented as a plain core.call and only diverges
       33 # under compilation (see named_call_translation_rule)
  
  AttributeError: module 'jax.core' has no attribute 'CallPrimitive'
  ```
  
  I believe that's because Haiku now requires `jax>=0.1.71`, while the version by default on Colab is `jax==0.1.69`. `CallPrimitive` was introduced in jax 0.1.71.
  https://github.com/google/jax/blob/1545a29e6d69a7b3c7fdf9a49b38004759a9fbfa/jax/core.py#L1106-L1115
  
  To reproduce (inside a Colab):
  ```python
  import jax
  print(jax.__version__)  # 0.1.69
  
  !pip install -q git+https://github.com/deepmind/dm-haiku
  import haiku as hk
  ```
  
  Run `!pip install -q --upgrade jax jaxlib` first in your Colab to fix this issue.
  
  </issue>
  <code>
  [start of setup.py]
  1 # Lint as: python3
  2 # Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
  3 #
  4 # Licensed under the Apache License, Version 2.0 (the "License");
  5 # you may not use this file except in compliance with the License.
  6 # You may obtain a copy of the License at
  7 #
  8 #     http://www.apache.org/licenses/LICENSE-2.0
  9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15 # ==============================================================================
  16 """Setup for pip package."""
  17 
  18 from setuptools import find_namespace_packages
  19 from setuptools import setup
  20 
  21 
  22 def _get_version():
  23   with open('haiku/__init__.py') as fp:
  24     for line in fp:
  25       if line.startswith('__version__'):
  26         g = {}
  27         exec(line, g)  # pylint: disable=exec-used
  28         return g['__version__']
  29     raise ValueError('`__version__` not defined in `haiku/__init__.py`')
  30 
  31 
  32 def _parse_requirements(requirements_txt_path):
  33   with open(requirements_txt_path) as fp:
  34     return fp.read().splitlines()
  35 
  36 
  37 _VERSION = _get_version()
  38 
  39 EXTRA_PACKAGES = {
  40     'jax': ['jax>=0.1.55'],
  41     'jaxlib': ['jaxlib>=0.1.37'],
  42 }
  43 
  44 setup(
  45     name='dm-haiku',
  46     version=_VERSION,
  47     url='https://github.com/deepmind/dm-haiku',
  48     license='Apache 2.0',
  49     author='DeepMind',
  50     description='Haiku is a library for building neural networks in JAX.',
  51     long_description=open('README.md').read(),
  52     long_description_content_type='text/markdown',
  53     author_email='haiku-dev-os@google.com',
  54     # Contained modules and scripts.
  55     packages=find_namespace_packages(exclude=['*_test.py']),
  56     install_requires=_parse_requirements('requirements.txt'),
  57     extras_require=EXTRA_PACKAGES,
  58     tests_require=_parse_requirements('requirements-test.txt'),
  59     requires_python='>=3.6',
  60     include_package_data=True,
  61     zip_safe=False,
  62     # PyPI package information.
  63     classifiers=[
  64         'Development Status :: 4 - Beta',
  65         'Intended Audience :: Developers',
  66         'Intended Audience :: Education',
  67         'Intended Audience :: Science/Research',
  68         'License :: OSI Approved :: Apache Software License',
  69         'Programming Language :: Python :: 3',
  70         'Programming Language :: Python :: 3.6',
  71         'Programming Language :: Python :: 3.7',
  72         'Topic :: Scientific/Engineering :: Mathematics',
  73         'Topic :: Software Development :: Libraries :: Python Modules',
  74         'Topic :: Software Development :: Libraries',
  75     ],
  76 )
  77
  [end of setup.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: diff --git a/setup.py b/setup.py
  --- a/setup.py
  +++ b/setup.py
  @@ -37,8 +37,8 @@
   _VERSION = _get_version()
   
   EXTRA_PACKAGES = {
  -    'jax': ['jax>=0.1.55'],
  -    'jaxlib': ['jaxlib>=0.1.37'],
  +    'jax': ['jax>=0.1.71'],
  +    'jaxlib': ['jaxlib>=0.1.49'],
   }
   
   setup(
  
verification_info: {"golden_diff": "diff --git a/setup.py b/setup.py\n--- a/setup.py\n+++ b/setup.py\n@@ -37,8 +37,8 @@\n _VERSION = _get_version()\n \n EXTRA_PACKAGES = {\n-    'jax': ['jax>=0.1.55'],\n-    'jaxlib': ['jaxlib>=0.1.37'],\n+    'jax': ['jax>=0.1.71'],\n+    'jaxlib': ['jaxlib>=0.1.49'],\n }\n \n setup(\n", "issue": "Jax version upgrade (AttributeError: CallPrimitive)\nUsing the current version of master 66f9c69 of Haiku, I am getting the following error on Colab\r\n```\r\nAttributeError                            Traceback (most recent call last)\r\n<ipython-input-3-3a9e6adbfff5> in <module>()\r\n----> 1 import haiku as hk\r\n\r\n/usr/local/lib/python3.6/dist-packages/haiku/__init__.py in <module>()\r\n     17 \r\n     18 from haiku import data_structures\r\n---> 19 from haiku import experimental\r\n     20 from haiku import initializers\r\n     21 from haiku import nets\r\n\r\n/usr/local/lib/python3.6/dist-packages/haiku/experimental.py in <module>()\r\n     22 from haiku._src.base import custom_getter\r\n     23 from haiku._src.base import ParamContext\r\n---> 24 from haiku._src.dot import to_dot\r\n     25 from haiku._src.lift import lift\r\n     26 from haiku._src.module import profiler_name_scopes\r\n\r\n/usr/local/lib/python3.6/dist-packages/haiku/_src/dot.py in <module>()\r\n     23 \r\n     24 from haiku._src import data_structures\r\n---> 25 from haiku._src import module\r\n     26 from haiku._src import utils\r\n     27 import jax\r\n\r\n/usr/local/lib/python3.6/dist-packages/haiku/_src/module.py in <module>()\r\n     26 from haiku._src import base\r\n     27 from haiku._src import data_structures\r\n---> 28 from haiku._src import named_call\r\n     29 from haiku._src import utils\r\n     30 import jax.numpy as jnp\r\n\r\n/usr/local/lib/python3.6/dist-packages/haiku/_src/named_call.py in <module>()\r\n     29 \r\n     30 # Registering named call as a primitive\r\n---> 31 named_call_p = core.CallPrimitive('named_call')\r\n     32 # named_call is implemented as a plain core.call and only diverges\r\n     33 # under compilation (see named_call_translation_rule)\r\n\r\nAttributeError: module 'jax.core' has no attribute 'CallPrimitive'\r\n```\r\n\r\nI believe that's because Haiku now requires `jax>=0.1.71`, while the version by default on Colab is `jax==0.1.69`. `CallPrimitive` was introduced in jax 0.1.71.\r\nhttps://github.com/google/jax/blob/1545a29e6d69a7b3c7fdf9a49b38004759a9fbfa/jax/core.py#L1106-L1115\r\n\r\nTo reproduce (inside a Colab):\r\n```python\r\nimport jax\r\nprint(jax.__version__)  # 0.1.69\r\n\r\n!pip install -q git+https://github.com/deepmind/dm-haiku\r\nimport haiku as hk\r\n```\r\n\r\nRun `!pip install -q --upgrade jax jaxlib` first in your Colab to fix this issue.\n", "before_files": [{"content": "# Lint as: python3\n# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.\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 for pip package.\"\"\"\n\nfrom setuptools import find_namespace_packages\nfrom setuptools import setup\n\n\ndef _get_version():\n  with open('haiku/__init__.py') as fp:\n    for line in fp:\n      if line.startswith('__version__'):\n        g = {}\n        exec(line, g)  # pylint: disable=exec-used\n        return g['__version__']\n    raise ValueError('`__version__` not defined in `haiku/__init__.py`')\n\n\ndef _parse_requirements(requirements_txt_path):\n  with open(requirements_txt_path) as fp:\n    return fp.read().splitlines()\n\n\n_VERSION = _get_version()\n\nEXTRA_PACKAGES = {\n    'jax': ['jax>=0.1.55'],\n    'jaxlib': ['jaxlib>=0.1.37'],\n}\n\nsetup(\n    name='dm-haiku',\n    version=_VERSION,\n    url='https://github.com/deepmind/dm-haiku',\n    license='Apache 2.0',\n    author='DeepMind',\n    description='Haiku is a library for building neural networks in JAX.',\n    long_description=open('README.md').read(),\n    long_description_content_type='text/markdown',\n    author_email='haiku-dev-os@google.com',\n    # Contained modules and scripts.\n    packages=find_namespace_packages(exclude=['*_test.py']),\n    install_requires=_parse_requirements('requirements.txt'),\n    extras_require=EXTRA_PACKAGES,\n    tests_require=_parse_requirements('requirements-test.txt'),\n    requires_python='>=3.6',\n    include_package_data=True,\n    zip_safe=False,\n    # PyPI package information.\n    classifiers=[\n        'Development Status :: 4 - Beta',\n        'Intended Audience :: Developers',\n        'Intended Audience :: Education',\n        'Intended Audience :: Science/Research',\n        'License :: OSI Approved :: Apache Software License',\n        'Programming Language :: Python :: 3',\n        'Programming Language :: Python :: 3.6',\n        'Programming Language :: Python :: 3.7',\n        'Topic :: Scientific/Engineering :: Mathematics',\n        'Topic :: Software Development :: Libraries :: Python Modules',\n        'Topic :: Software Development :: Libraries',\n    ],\n)\n", "path": "setup.py"}]}