Add __version__ attribute to NT

PiperOrigin-RevId: 318825275

docs/conf.py

@@ -12,6 +12,10 @@
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 #
+
+"""Readthedocs configuration."""
+
+
 import os
 import sys
 sys.path.insert(0, os.path.abspath('..'))
@@ -22,10 +26,6 @@
 copyright = u'2019, Google LLC.'
 author = u'The Neural Tangents Authors'
 
-# The short X.Y version
-version = u''
-# The full version, including alpha/beta/rc tags
-release = u'0.3.0'
 
 # -- General configuration ---------------------------------------------------
 
@@ -194,4 +194,4 @@
 # -- Options for intersphinx extension ---------------------------------------
 
 autodoc_mock_imports = [
-    "frozendict", "jax", "absl", "numpy", "scipy"]
+    'frozendict', 'jax', 'absl', 'numpy', 'scipy']

neural_tangents/__init__.py

@@ -15,6 +15,10 @@
 
 """Public Neural Tangents modules and functions."""
 
+
+__version__ = '0.3.1'
+
+
 from neural_tangents import predict
 from neural_tangents import stax
 from neural_tangents.utils.batch import batch

setup.py

@@ -38,9 +38,29 @@
   INSTALL_REQUIRES += ['dataclasses>=0.7']
 
 
+def _get_version() -> str:
+  """Returns the package version.
+
+  Adapted from:
+  https://github.com/deepmind/dm-haiku/blob/d4807e77b0b03c41467e24a247bed9d1897d336c/setup.py#L22
+
+  Returns:
+    Version number.
+  """
+  path = 'neural_tangents/__init__.py'
+  version = '__version__'
+  with open(path) as fp:
+    for line in fp:
+      if line.startswith(version):
+        g = {}
+        exec(line, g)  # pylint: disable=exec-used
+        return g[version]  # pytype: disable=key-error
+    raise ValueError(f'`{version}` not defined in `{path}`.')
+
+
 setuptools.setup(
     name='neural-tangents',
-    version='0.3.0',
+    version=_get_version(),
     license='Apache 2.0',
     author='Google',
     author_email='neural-tangents-dev@google.com',