add jax to setup, set 0.1 version

Author: parrt

setup.py

@@ -25,7 +25,8 @@
 
 tensorflow_requires = ['tensorflow']
 torch_requires = ['torch']
-all_requires = tensorflow_requires + torch_requires
+jax_requires = ['jax', 'jaxlib']
+all_requires = tensorflow_requires + torch_requires + jax_requires
 
 exec(open('tsensor/version.py').read())
 setup(
@@ -40,8 +41,9 @@
     install_requires=['graphviz>=0.14.1','numpy','IPython', 'matplotlib'],
     extras_require = {'all': all_requires,
                       'torch': torch_requires,
-                      'tensorflow': tensorflow_requires
-                     },
+                      'tensorflow': tensorflow_requires,
+                      'jax': jax_requires
+                      },
     description='The goal of this library is to generate more helpful exception messages for numpy/pytorch tensor algebra expressions.',
 #    keywords='visualization data structures',
     classifiers=['License :: OSI Approved :: MIT License',

tsensor/version.py

@@ -21,4 +21,4 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 """
-__version__ = '0.1b7'
+__version__ = '0.1'