CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!
CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!
| Download
Project: Testing 18.04
Path: jax.ipynb
Views: 877Kernel: Python 3 (Ubuntu Linux)
JAX on CoCalc
Kernel: Python 3 (Ubuntu Linux)
JAX is Autograd and XLA, brought together for high-performance machine learning research
grad for gradient, and jit for just in time compilation
jax also has a numpy compatible interface
In [1]:
In [2]:
In [3]:
0.4199743
/usr/local/lib/python3.6/dist-packages/jax/lib/xla_bridge.py:164: UserWarning: No GPU found, falling back to CPU.
warnings.warn('No GPU found, falling back to CPU.')
In [4]:
4.58 ms ± 551 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [5]:
In [6]:
In [7]:
In [8]:
146 ms ± 2.95 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
In [9]:
466 ms ± 5.45 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
In [10]:
comparing with numba
In [11]:
In [12]:
In [13]:
612 ms ± 22.5 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)
In [0]:
In [0]: