Version bump to make NT work on M1 macs (#171). BREAKING: drop suppor…

…t for Python 3.7 (necessary due to JAX and TF2JAX policy). Also fix some linter errors. PiperOrigin-RevId: 508726120
google · Feb 10, 2023 · 0272225 · 0272225
1 parent 33fc4de
commit 0272225
Show file tree

Hide file tree

Showing 5 changed files with 63 additions and 39 deletions.
diff --git a/.github/workflows/linux.yml b/.github/workflows/linux.yml
@@ -21,7 +21,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: [3.7, 3.8, 3.9, '3.10']  # 3 == 3.11 currently.
+        python-version: [3.8, 3.9, '3.10']  # 3 == 3.11 currently.
         JAX_ENABLE_X64: [0, 1]
 
     runs-on: ubuntu-latest
@@ -32,10 +32,10 @@ jobs:
       uses: styfle/cancel-workflow-action@0.11.0
       with:
         access_token: ${{ github.token }}
-    - uses: actions/checkout@v3.2.0
+    - uses: actions/checkout@v3.3.0
 
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4.3.1
+      uses: actions/setup-python@v4.5.0
       with:
         python-version: ${{ matrix.python-version }}
 

diff --git a/.github/workflows/macos.yml b/.github/workflows/macos.yml
@@ -21,7 +21,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: [3.7, 3.8, 3.9, '3.10']  # 3 == 3.11 currently.
+        python-version: [3.8, 3.9, '3.10']  # 3 == 3.11 currently.
         JAX_ENABLE_X64: [0]
 
     runs-on: macos-latest
@@ -32,10 +32,10 @@ jobs:
       uses: styfle/cancel-workflow-action@0.11.0
       with:
         access_token: ${{ github.token }}
-    - uses: actions/checkout@v3.2.0
+    - uses: actions/checkout@v3.3.0
 
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4.3.1
+      uses: actions/setup-python@v4.5.0
       with:
         python-version: ${{ matrix.python-version }}
 

diff --git a/.github/workflows/pytype.yml b/.github/workflows/pytype.yml
@@ -27,12 +27,12 @@ jobs:
       uses: styfle/cancel-workflow-action@0.11.0
       with:
         access_token: ${{ github.token }}
-    - uses: actions/checkout@v3.2.0
+    - uses: actions/checkout@v3.3.0
 
-    - name: Set up Python 3.7
-      uses: actions/setup-python@v4.3.1
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v4.5.0
       with:
-        python-version: 3.7
+        python-version: '3.10'
 
     - name: Install dependencies
       run: |

diff --git a/.github/workflows/sketching.yml b/.github/workflows/sketching.yml
@@ -20,7 +20,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: [3.7]
+        python-version: [3.10]
         JAX_ENABLE_X64: [0, 1]
 
     runs-on: ubuntu-latest
@@ -31,10 +31,10 @@ jobs:
       uses: styfle/cancel-workflow-action@0.11.0
       with:
         access_token: ${{ github.token }}
-    - uses: actions/checkout@v3.2.0
+    - uses: actions/checkout@v3.3.0
 
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4.3.1
+      uses: actions/setup-python@v4.5.0
       with:
         python-version: ${{ matrix.python-version }}
 

diff --git a/setup.py b/setup.py
@@ -26,10 +26,10 @@
 
 
 INSTALL_REQUIRES = [
-    'jax>=0.3.13',
+    'jax>=0.4.3',
     'frozendict>=2.3',
     'typing_extensions>=4.0.1',
-    'tf2jax>=0.3.0',
+    'tf2jax>=0.3.3',
 ]
 
 
@@ -44,20 +44,20 @@ def _get_version() -> str:
   """Returns the package version.
 
   Adapted from:
-  https://github.com/deepmind/dm-haiku/blob/d4807e77b0b03c41467e24a247bed9d1897d336c/setup.py#L22
+  https://github.com/deepmind/tf2jax/blob/41ea640c7525ed42f7dcd02937b1f308f8949521/setup.py#L24
 
   Returns:
     Version number.
   """
-  path = 'neural_tangents/__init__.py'
-  version = '__version__'
-  with open(path) as fp:
+  current_dir = os.path.dirname(os.path.abspath(__file__))
+
+  with open(os.path.join(current_dir, 'neural_tangents', '__init__.py')) as fp:
     for line in fp:
-      if line.startswith(version):
-        g = {}
-        exec(line, g)
-        return g[version]
-    raise ValueError(f'`{version}` not defined in `{path}`.')
+      if line.startswith('__version__') and '=' in line:
+        version = line[line.find('=') + 1:].strip(' \'"\n')
+        if version:
+          return version
+    raise ValueError('`__version__` not found in `neural_tangents/__init__.py`')
 
 
 setuptools.setup(
@@ -73,27 +73,51 @@ def _get_version() -> str:
     url='https://github.com/google/neural-tangents',
     download_url='https://pypi.org/project/neural-tangents/',
     project_urls={
-        'Source Code': 'https://github.com/google/neural-tangents',
-        'Paper': 'https://arxiv.org/abs/1912.02803',
-        'Finite Width NTK paper': 'https://arxiv.org/abs/2206.08720',
-        'Video': 'https://iclr.cc/virtual_2020/poster_SklD9yrFPS.html',
-        'Finite Width NTK video': 'https://youtu.be/8MWOhYg89fY?t=10984',
-        'Documentation': 'https://neural-tangents.readthedocs.io/en/latest/?badge=latest',
-        'Bug Tracker': 'https://github.com/google/neural-tangents/issues',
-        'Release Notes': 'https://github.com/google/neural-tangents/releases',
-        'PyPi': 'https://pypi.org/project/neural-tangents/',
-        'Linux Tests': 'https://github.com/google/neural-tangents/actions/workflows/linux.yml',
-        'macOS Tests': 'https://github.com/google/neural-tangents/actions/workflows/macos.yml',
-        'Pytype': 'https://github.com/google/neural-tangents/actions/workflows/pytype.yml',
-        'Coverage': 'https://app.codecov.io/gh/google/neural-tangents'
+        'Source Code':
+            'https://github.com/google/neural-tangents',
+
+        'Paper':
+            'https://arxiv.org/abs/1912.02803',
+
+        'Finite Width NTK paper':
+            'https://arxiv.org/abs/2206.08720',
+
+        'Video':
+            'https://iclr.cc/virtual_2020/poster_SklD9yrFPS.html',
+
+        'Finite Width NTK video':
+            'https://youtu.be/8MWOhYg89fY?t=10984',
+
+        'Documentation':
+            'https://neural-tangents.readthedocs.io/en/latest/?badge=latest',
+
+        'Bug Tracker':
+            'https://github.com/google/neural-tangents/issues',
+
+        'Release Notes':
+            'https://github.com/google/neural-tangents/releases',
+
+        'PyPi':
+            'https://pypi.org/project/neural-tangents/',
+
+        'Linux Tests':
+            'https://github.com/google/neural-tangents/actions/workflows/linux.yml',
+
+        'macOS Tests':
+            'https://github.com/google/neural-tangents/actions/workflows/macos.yml',
+
+        'Pytype':
+            'https://github.com/google/neural-tangents/actions/workflows/pytype.yml',
+
+        'Coverage':
+            'https://app.codecov.io/gh/google/neural-tangents'
     },
     packages=setuptools.find_packages(exclude=('presentation',)),
     long_description=long_description,
     long_description_content_type='text/markdown',
     description='Fast and Easy Infinite Neural Networks in Python',
-    python_requires='>=3.7',
+    python_requires='>=3.8',
     classifiers=[
-        'Programming Language :: Python :: 3.7',
         'Programming Language :: Python :: 3.8',
         'Programming Language :: Python :: 3.9',
         'Programming Language :: Python :: 3.10',