<jupyter_start><jupyter_text>DESC JAX TutorialAuthor:- [@EiffL](https://github.com/EiffL)(<NAME>) OverviewThe goal of this tutorial is to cover some of the essential JAX mechanisms, andput them into practice for typical computations that may be required for a project like CCL**Learning objectives**:- Learn about JAX function transformations- Learn about common gotchas when using JAX- Implement integrations and ODE in JAX  JAX 101We will start by a brief overview of JAX fundamentals, this is an abrieged version of what you can find in the official JAX documentation:- [JAX Quickstart](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html)- [JAX 101](https://jax.readthedocs.io/en/latest/jax-101/index.html)<jupyter_code>%pylab inline
import jax
import jax.numpy as np
import numpy as onp    # We'll import good old numpy as onp for reference if we need it<jupyter_output>Populating the interactive namespace from numpy and matplotlib<jupyter_text>Random Numbers (let's get it out of the way)JAX has a much more evolved approach to random number generation than NumPy, design to allow parallel random number generators. It is also needed to statisfy the JAX pure function approach.To use any stochastic function in JAX, you need to specify a key, which is a seed that the function can then use internally. So let's create one:<jupyter_code>rng_key = jax.random.PRNGKey(42) # Because 42 is the answer<jupyter_output><empty_output><jupyter_text>Then everytime I need a key for a stochastic operation I can use:<jupyter_code>key, rng_key = jax.random.split(rng_key)
random_array = jax.random.normal(key, shape=[10])
key2, rng_key = jax.random.split(rng_key)
random_array2 = jax.random.normal(key2, shape=[10])
random_array2, random_array<jupyter_output><empty_output><jupyter_text>this key is only used for this operation, and next time I need to run an stochastic operation I can split again `rng_key` to obtain a new key.  JAX works just like NumPy (most of the time)Let's start with very simple stuff, a matrix multiplication:<jupyter_code>key1, key2, rng_key = jax.random.split(rng_key, 3)

# I'm gonna create some random matrices
A = jax.random.normal(key1, [500,1000])
B = jax.random.normal(key2, [1000, 500])
C = np.dot(A, B) 
print("Shape of result:", C.shape)<jupyter_output>Shape of result: (500, 500)<jupyter_text>Under the hood, if you have a GPU, this matrix multiplication happened on GPU ;-)<jupyter_code>C.argmax() # Random numpy operation<jupyter_output><empty_output><jupyter_text>But you also have access to a lot of other linear algebra tools:- [jax.scipy.linalg](https://jax.readthedocs.io/en/latest/jax.scipy.htmlmodule-jax.scipy.linalg)- [jax.scipy.optimize](https://jax.readthedocs.io/en/latest/jax.scipy.htmlmodule-jax.scipy.optimize)- [jax.scipy.special](https://jax.readthedocs.io/en/latest/jax.scipy.htmlmodule-jax.scipy.special)<jupyter_code>from jax.scipy.linalg import svd
U, s, V = svd(C)  # Example of an SVD
from scipy.linalg import svd
U, s, V = svd(C)
U<jupyter_output><empty_output><jupyter_text>Just in Time CompilationWhen you execute JAX code without JIT, you run through the code at the Python level, until you encounter the low level XLA interface, which is hidden behind the numpy API. At that point, the XLA bit of code is executed, and the result is returned to Python which continues to run through the next commands.This can be pretty slow because the execution is still driven by Python.So the idea of JAX is that you can jit a big function, to turn it into a single,compiled XLA graph, that runs without needing Python<jupyter_code>def func(A, B):
  C = np.dot(A, B)
  C = np.where(C>0, C, 0)
  return C
%timeit func(A, B).block_until_ready()
jitted_func = jax.jit(func)
%time jitted_func(A, B).block_until_ready(); # First execution won't be fast
%timeit jitted_func(A,B).block_until_ready() #  Next calls are fast<jupyter_output>The slowest run took 8.99 times longer than the fastest. This could mean that an intermediate result is being cached.
10000 loops, best of 5: 138 µs per loop<jupyter_text>For fun, we can compare to CPU using good old numpy:<jupyter_code>A_np = onp.array(A)
B_np = onp.array(B)

def func_onp(A, B):
  C = onp.dot(A, B)
  C = onp.where(C>0, C, 0)
  return C
%timeit func_onp(A_np, B_np)<jupyter_output>100 loops, best of 5: 6.13 ms per loop<jupyter_text>Auto-batching with vmapAny function can be "batched" which means that you add a new dimension to the input/output tensors. So, let's say I want to make my previous function compute that operation for 10 matrices at once.If I try to use the function directly, it won't work, because  that function was written thinking about single input matrices:<jupyter_code>A2 = jax.random.normal(key1, shape=[10, 500,1000])
B2 = jax.random.normal(key2, shape=[10, 1000, 500])
func(A2, B2).shape<jupyter_output><empty_output><jupyter_text>But we would have expected an output of shape [10, 500, 500].So, let's use vmap:<jupyter_code>batched_func = jax.vmap(func)
batched_func(A2,B2).shape<jupyter_output><empty_output><jupyter_text>Boom :-) We just turned a function designed to work on a single example of matrices into a batched function.And you can also compose jit and vmap<jupyter_code>cool_func = jax.jit(jax.vmap(func))
cool_func(A2,B2)
%timeit cool_func(A2,B2).block_until_ready()<jupyter_output>1000 loops, best of 5: 780 µs per loop<jupyter_text>Note that this takes less than 10 times the execution time of a single instance.  Getting gradients with grad, jacfwdObviously, one of the things we care about the most here are gradients, so let's give it a try.<jupyter_code># I'm gonna create an array of some parameters
params = jax.random.normal(key1, shape=[500])

def f(params, A, B):
  C = np.dot(A, B)
  C = np.where(C>0, C, 0)
  return np.sum(C.dot(params)**2)
f(params, A, B) # Ok so this function does something, ok:<jupyter_output><empty_output><jupyter_text>Let me take the derivative of this function. jax.grad will compute the derivative of this function with respect to the first argument of f:<jupyter_code>df_dx = jax.grad(f)
df_dx
g = df_dx(params, A, B)
g.shape # And we can check that the grad has the same dimensionality as the
        # input array<jupyter_output><empty_output><jupyter_text>But what happens if I want to take the Jacobian of a  non scalar function? Well, let's just use jax.jacobian<jupyter_code>def f2(params, A, B):
  C = np.dot(A, B)
  return C.dot(params**2)
f2(params, A, B).shape
df2_dx = jax.jacobian(f2)
df2_dx(params, A, B).shape<jupyter_output><empty_output><jupyter_text>Boom, we have computed:![image.png](data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAe8AAACJCAYAAAAFU7A+AAAABHNCSVQICAgIfAhkiAAAIABJREFUeAHtncGLLFmb1r8YhruRcbI+cSUDmYgy6KZKB3fSEY3MYjZflYjgIH6ZDW5GsKvGjSBMR/wB0nkXwyxcdCY4HwjCzbtSEKxz/wDt6o0rsfKCuHFxiwHRXfv8qs+ZPh03IjMiMiIzbuX7wFPnnPec855znjfyvJlZdbt/9jODKWAKmAKmgClgCpgCpoApYAqYAi9MgUTnqeILO+bZH4cYnwJVz1awnWI/tqYpYAqYAp+8AoVO8H0N2172E/l5IzrPb1VeiftwrQHsg7n4MPSnwI1cOc8HlUtxn8b05+JKnIuHwslB1TNGzA2mgClgCpgCHRQIyZuyzLbunCbEyfpvqf2hZFPzJ2D8k3grcsF/Lhr6USCVm03J1Vrt+5Kt3GTMUuTN11Y8FAs5KD9bxBqbwRQwBUwBU6CDAlygXKRJh7nxlGvv5y42qv5WLCeQeMhajXfijR83iTutfpAC6E5sY1yogS2LjVF9FvUzfx719VVN5MiSd19qmh9TwBQ4SwUKnbqP5H0pP1z2UzFG8B/b4rpd4rEa/dZzuVtWuETzdYUd061IP0l+KCRybHEfSl3zawqYAmehQEiuXKh94EpOSACpd7bLf6YxXOLXfqwV/StAEkZfYjITiTOaO7EKaxm3VR092sIeih59mitTwBQwBc5KAS5QLnMu1EMw02R+R7oWM5FkUYhOLPsnwd+LjKePknYqGvpTYCFXW/FOJCYbERuavxZjLNUgBh88qTN+CCRyyh6KIZybT1PAFDAFzkEBLlAuUi7UrphpIpc+SSIGCQL7U2yM6iQM1v418ZD1I5dW9QoQV3SfRIpQJx51iZMYhL4h4xGvE23PqqaAKWAKmAJNFeCSPzR538vHtmLBXDZ8ryv6MDlP6jFIMl+Jl7HR6o0VuNJIdC+/mcKBE+ljTBlh3k25Q21ikVfYu5gseXdRzeaYAqaAKRApcGjyzuSLZPA68hmqwXdVMggX+DIM9iXtXMRnKhraK7DRFPSbVUzFvq2wY1qIVfNWsjtPFQcjxL442JM5MAVMAVPgTBXgAuXC5kLtgjC/KkF/K4fvxSrfV7Kz7kKsAn1pVYfZdiqA1o/iVizrzqdndP1CrAJvwJ7E8jzG5qIT+wD+2UfRhzPzYQqYAqbAOSrABcpFWnVhN9EjzJ+VBtPGb1VSZ+jC95PEq8DctKrDbHsVQLt3FaNWspGcJxV9mJwn9TJyGVzZ2LGdaB57LDrOt2mmgClgCpy9AlygXKRcqF2w0CTmZ9FkksODuI5s5Wohw6516UvLk6zdSAGnUXzrEYP4oGndm6XE9xOXKuQyuqqODrZ9a3VwaVNMAVPAFDgvBbisdyXRJmqsNYhkMfd8VFmXBNT1DKefsA7sKavrNPtOBS7VuxWX4me+5M3ULj1J6mhe901Jrj4n9oFETlhr3zPSx1rmwxQwBUyBF6kAFygXKRfqIcg0GV8LcSruQqJO/skS4+vAnrK6TrPvVWCiEQsRjUnIF+IuMBbN68bl6nNiH0jkxJJ3H0qaD1PAFDhbBbjc+0jeTQS81qA34lRkzZlYB/qzuk6z96bAN/LEH7DxKd2JdcjV4eo6W9oTjSe+Rct5NtwUMAVMAVPAK8AFykXKhTo0nBZ4Em/FjViFuYz3Inviq3iSimEYBTK5RWc0fhRpV4F40M+3JdTn4iHgWfteLA5xYnNNAVPAFDhnBbhAuUi5UIfGXAuQtEkWE9FwWgWIAfGAiyNuhWfNkvcRBbelTAFT4OUpcMzk/fLUsxN1UcCSdxfVbI4pYAqYApEClrwjMax6FAUseR9FZlvEFDAFXrIClrxfcnTHeTZL3uOMi+3KFDAFPiEFLHnXB+uyvuu5Z6qf/N64DvRN6zrP2G7J+4yDb0c3BUyBfhTokrx/W0v/2xPyH/Zz9J1e5urlj6pWYhWmMvLX17AO9+rAx743AXXzm9p/RwNPGQ/W/qOmm9U4S94txLKhpoApYApUKdAlef+WHP3LE/KzqoP0bLuWPxLvssYvn6q34kNNP+aV+CROxSHx1+X8lPFg7X/c4oCWvFuIZUNNAVPAFKhSoEvyrvJjNlOgqQKWvJsqZeNMAVPAFKhRwJJ3jTBmHkwBS96DSWuOTYF+FOCrxayG9PUNfKY7eNX3gjX+ZrJXnftY69dsq9JcyMrXw0llrxlNgf4V4FnjmePZM5gCpsAIFci0J16kVUwH2O9CPh/EW/FaZN2VmIpL0YnHAJdS1Zmd7MkxNtBijbDXvvfVt78WR7KhFQqMKR7shddHUbFPM5kCpsAIFMi0B16kd+LnJU7U7hu5HE4jp6ydR+24Hpl7r87ksXzeB9mcmIhjQqHNoFMf+wpvkN7J33firWg4jQK8vr4Rw3NHPG5Os5WPVuVZ45nj2TOYAqbACBXItCdepJRNMdNAxrcl/uf8iMDaedRumkzYw1Lk8uMS7ANOTmAijgmFNoNOh+4LvXiTFnChylZcBIOVR1XgTUn7mdpP4hgSOM8az1whGkwBU2CECmTaEy9SyqYIc5jXlO9qnDM/r+mrM5N0+Le7c18WKvuAkxOYiGMC50OnQ/Y19T42KmOQzHf9O+h4rNX7U+BSrojpquSSWD+WbKdo8qyxP/ZjMAVMgREqkGlPvEgp28BpcHhxMzd8BR3XsdF+50sVHwEf+UfW3YYv1L0V+YTyIHIR9gEnJzARx4RCm0GnQ/Y10fyNmIoxMjW6xD/2YfX2Ckw1xYnlZ5dnmnhciacEzxr74NkzmAKmwAgV6Hp5h0vmfs+ZZup3O8ZwQeQ7+qu6vpWRRNQ3nBxCLq4xgQsUnfrYF/GYe5LQQxwz1Q3HV4AkPfdk9RDrU8cj0V4seRMRgykwUgW4JHiRdrkstn7uTGUd1urY5btt8k7kb6hLxck3ZI0x4T9oM5z51w/YFImaNzxOJGHDlViIcfxJJtideC9OxTDva9UN/ShwITfoi7YhHk7112I5HkH/EI+VxjgxF4cCzxr7YM2xvR6GOrP5NQU+KQUy7Ta+LNpsfuHnrmsmzWR3NX2YJyJr5+I+TDWAi4RP3cx59O25yr7g5Agm4pjwb7QZzvwbHTeFzlXfVpA08AsvRRKKExkPnEgfcVz5ugrDgQqEeCxLfhZqh3jQhe4bMY7Hk9rE6V5k7FB4Jcf/R/z34theD0Od2fyaAp+UApl2yyVA2Ra8qN+LzJ9VTF7LtqiwY8rFIiqpT8VdyNTJhcd6f1+kHS42VQ+GkweYiGPCv9ZmOHPX5F34+bPSoVJvf6+SM/OpLxMD4oS/khHtY1yq8U1ssHojBdCZeJKEY6RqYH/wxpXKOGaPaq9930Zl7utDFCTvPxNXIs+GwRQwBUamAJc1FwZlFyw0ifnr0uSZ2tuSrY8mF9+TmFQ4S2UjoXSF00RY5Vvmk+GQ5M1Z4oQQH+La96FpGRMZmPeH5Q7fXqmEjDE0V4DXBZptKqYsfN9dRd/U9zHmGLDkfQyVbQ1T4AAFMs3lMqHsgkST3ov4mEUO1qovonZfVSdHMEaqxkrciqnYFU4TYSKOCYck70wHITZVCXrp+2YVh73xfVcVfcGU+jGhbeV+BUI8ioqhK9mexIuKvoVs5ddYxbDeTJa8e5PSHJkCwygQLhPKrggXy9o7mKl8LyZdHdbMwx8XWFHT72RPa/qamJ0Gwb73LZcHoY/kvajYwQfZ4qSeqn3rx4XEHrSYyJ77vlCkqhAPQ3MF7jQUzbLSFPQlHkVkT1W/8e2Vyq0Y4jFVPcSKOvH6WmR8Lq7EqdgVlry7KmfzTIEjKcAlUnWZtFmeC2Xr/UxVrsWF2Dcu5ZC9hgut7N/JkJaNLdpOYyHnGRMOSd4zHQTNSBoxVmo8iJPISPJ48u1Hlc7XKQoxi9pUUxHfhuYK7IqHK7kJ8SBG1OP+ldohHkvVpyKxWIgg96TeBZa8u6hmc0yBIyrABcCLPlwEXZfm0sDPvfheTMS+EdaY1Th2sqc1fU3MToPgEHuX2844JHmzKLpx+ZPAfyE68a0YJ241n+1fqXwjFuJW/FJcibTLSGUg5oZ2CoR43GratchrpioeD7J/La7EL8StSDw2Iu2AqSr4dGKAU2UeGh3KV5pjf7DWQTibYgocS4FMC3EBUx6CRJO3Ir64SIZAIadPImtVwcmYVXU0tDEf1vlX10lwaPJm01ciyRtSr8JERvSb+c5yuzwnlcGSd1mVZu028bjwLimJz9S342KlRuENxI24UHaFJe+uytk8U+BICnAZ9JG82e61uKEyEJz8vt3hm37O0xVOE2Eijgl9JO8hzpPKqSXvIZRt7/NRU8KzP1ed1+FUnItdYMm7i2o2xxQ4ogK84PtK3kNs+1JOvxInfp+LHYs49YULbMew2i7mw0QcE8aYvJcSCDpfzlUaTqfAU7T0ter3IvHpiheTvHkxx+wqiM07TwXiZ2dsiSFTSMacvDd+f3OVW7EKqYwk+A/iG19X0RpOM+DYYjTG5C2ZDC9YgReRvP9IAXIlcuEZTIEmCvAHJq5E3hmPBWNP3mi1FZ14Ke7DIYnXyTk8xIem9w5L3r1Lag73KPBK/Z/8H6z9qQ6xEj+P+JdVPwQzTV6JfLVByQVaBf5AIRX5VHErTkTD4QqgN5/Q0H8pEo8qYL8W0f+XVQMa2P6axsTPzn9R+w8azDvWELT4XqQ8dzgJABNxTLDkPaZonMdeXkTy/pVilfcYr0v5cuKV98mluRUXvh0K2nwNuBWdyAV7LxoOU2Cu6WtxJiZi0DnEQ6ZnkNTR3IlbkXouHor/JAf/7FAnPc7P5IuzUZ47nASAPBdjgiXvMUXjPPbySsf85D95/0qHyHuMFwkYxuDTHYmaT9pgJnKhwolY+PqTSkN3BdASnfOSi9dqu8i2UB3tg42S9kY8FJa8D1VwuPlOrqElb4lgOGsFXun0lrxLjwAJ2JVsJG2Sw523Z77tfJukcyumvm1FNwUuNQ2d89L0kKxn3l6ojMcxD9tUPBSWvA9VcLj5Tq5hIo4J9sl7TNE4j71Y8t4RZy6I8EmbOsmi8OMz33a+bUW/CgTtKcGNiP7oDpZiHI9nY08/LHn3JOQAbpx8wvBcqDoKWPIeRRjOahOWvCvCfSWbEzfia1+G5FGo/YX4rUjy4Cte56nC0IMCC/lw4toTzWFI3iTuR9+mdCK2vmDJuy8l+/fj5BIm4phgyXtM0TiPvbzSMe1r8yjWC9WfRBJ4QKZKSBYkEdorkWQCsUHD4Qp8IxcP4kXkCm2D/sRlIToR7bci/dj6giXvvpTs34+TS5iIY4Il7zFF4zz2Ysk7ivNMdT5JkwxicFGQKGBI6llko99wuAI3coHGaclV0Jo3VUHrQnXGOrFvWPLuW9H+/Dm5guE5UHUUsOQ9ijCc1SYseUfh3qhOQphFNqpcFNjf+7qK50/f2ODYLhL29ymCN07bio1nsqHzOuorvM1Ftr6qlrz7UrJ/P04u4dhec5a8FRTDURWw5O3l5jLYlzxIGAGZKpa8gxqHl3yjgZ6vK1yhO31oHhBsLhh6LC159yhmz66c/MFEHBMseY8pGuexF0veUZxJEKuoHapLVfjK9iIYVJJILHlHghxYDV+Zzyv8PMrmSnZL3iVBzqTJcwAteUsEw1krYMk7Cj8Jeh21qV6KfCJf0IiQqW7JOxLkwGr45F3WmTZxoT9GoQb6u9jYU90+efck5ABunHzCRBwT7JP3mKJxHnux5B3FmURBog6JIlOdT313YoylGt+KIXk71fnnY4bDFFhr+r04EyfiXCRxX4oxnBrEBf2JlxPLY2TqDEvenaUbfKLTCjARxwRL3mOKxnnspVXy5oJ0IhdsuVzKdir8SgvnPS1+Iz9vRSeSTEIiV/XPMVetiJirnoqGwxXgjZLz5Jm6EMsI2ufqCPVpedAB7bEl73+ls/BG5V8ccKaXMvV/6yD/V0xGdiBL3iMLyBls55XO2PjfeU81OFyWTnUuFLgWF+Kp0GfyPtUZbN3xKDC25P0LScPr7O+dWKLJiddn+Tfi/xAteaOG4ZwVaJW8Y6FI4iF5Z3HHCeqWvE8g+gte0pL3x8G9lonX+ym/YWNXlrxRwWAK/OxnlrztKTAFSgpY8i4JoqYl7481iS32tXmshtWPoYAl72OobGt8UgpY8h5vuOyT9w+xGcubqfE+Kcff2b2W/CBOj7T0UZM3v6dqyjbnt6/N26hlY/cp0CV5N32uGdcWv9CEMfzOu+2+hxh/SPIeMkbH/uT9qSTviR6Cr8TLIR6Gkfl02g+v02Od9WjJu9CheGdSRSd7mTycTWHJu6lSNq6JAm2T95DPNvu15P1j1Lom741cVN092FwFp7K1wbGTd5u9nXJsrsW/F51o6FeBV3LX+K/N46W5sAgKzOKOHfVEfWUyvGyj3QZNk/dfktO/e0L+lTaHajF2orFcavmOOV/7MYytwlxGLrLLqs4ebH9RPk6pPWv/TsNztE3euK17huvsDbfyPMyS949qdU3eeKiLRZ39x1X31yx5V2vEfZOLQ90r1aueh7WX5J2eWKumyft3tc/HE/KfDqRTKr+8iXra4Z8+xtS9iDa+/3aHj0O6/rYmn1J71nYND9AleTd03WmYJe8fZTskef/opf+aJe/+NTWPuxVolLz5D5Xci3nkq1CdZACzyF5XxQfjmnJa56jC3jR5V0x9Maa5TjLdcZpUfdc7+ifqm4uU5462yXvIZ5tYjCF5T7UPEudcPCW6Ju9Um2569zCu7evAkvcpn4rzXLtR8nbSJiRqLipQiMHGw74PCw1gTlOm+xxG/Za8IzGserACbZP3kM82hxlD8r7VPni9O/GU6Jq877TppncP46YtD2nJu6VgNvxgBRol742W4YX7JE7Fn4uPIrZ34qlhyfvUEXhZ67dN3kOffgzJm0+iJPDp0Ifd479r8t7j9uDuYyfvqXaMFqn4KYPn6oP47ad8CL/3uUpiwpmOgUbJm81sRJL1vS+pF+KxNqqlamHJu1Ya6+iggCXvDqIdaQqXo/3nUX94I8Ud7I6k+1DLkD+exO1QCxzR74PWIibpkdZslLzDXhJVZiIlHAsOTd6ZDrIUeWPC12sXouE4CvA8fS2ifS6OQXtL3grESGHJ+4fAkPRuxcuRxqnNtjjLSwCxICbHQqvkfaxNtV3nkOQ912Ir8UokkRTioziGJKJtvGhc63QPIm+e0PtO/CASi1PCkvcp1d+9tiXv3fpY7/kocNbJe6I4kyzK/y/utWx8EjQMqwDa82Ypxms1XGw4QX1syfv3pcH34i9OoMXYlvzP2tD/FJORbezYv/Me2fFtOydQ4KyT96UE51LMS8LzCRD7Rcluzf4UCNovSy5v1Eb7q5L9mM2xJe9/4jX5R8cUoWKtaYXt2Kb/pgX/u2jJ+9jKn269VEt/ebrlR7vyWSfvOCpcBuFCWKhOAsniAVbvXQG++QCx9mh+au3HlrzHoMm1j8tK5SnhtDjkmRkT7JP3cNF4kGvuhHS4JT5Jz2efvPmktxGdJ1+hF2I5gdzKlouMB6F9+UPTfnZQgKR0LzrPO5WwSvuvZJ+LgBgsReYPAUveH6t6LRNxQfdTwmlxaMl7HP/SR6EYHDx7K3Ey+Eqf1gJnnby5iB7EWRSz16o/ilxUwe5UJ2Hw8GzFb8WFSKLh97aG9gqg3ZN4FU0tVEdbtM+8nRih9VTcivci47ANpb0lb4k7UjjtCybimHDsT97XOjyvE14fhnEowN1ETC6PtJ2zTd4kY4ROS0Jn3r5VyQVBkmFsAMnlwTeeonrop0zjhtU/UiBoPC/1BDtxQfuFiP4BTpWtb1BH/yFgyXsIVfvx6eQG8nyMCZa8xxSN0+xlo2W5k6ZHWv5skzef2rYVImeykTzWvu86GsOFQV/hbVOVE1+n+Ep0ImMM9Qo8qouHvIyg/VvfkaoM+qI9MXstgqkY+mj3CUvefarZry8nd5DnYUw4dvIe09ltL6dR4CyTd0gSIRHE0hdqkHwZU8aNDHV9YWzqx4S2lT9VYOb1Wf3U/Ny6831V2l/5PmIwNCx5D61wd/9OU2EijgmWvMcUjfPYy1km75AkqhIBnwpdTexfy07yjnEZN1RPxfKY0pCzbpKY0aeoUMHJtq2wYwoxiz9tx9r/UmNSBgpz8UsxHqtmY1jybizV0Qc6rQgteUsEw1krcJbJOySQcvKe61F4Eq+iR4LfcT/6NnXn6xQLz8hkyTsWo6KOtiRvtIuRqoE91v5ebb4qB2/F78RwaRPDQgRzMRWZvxTxQXkrdoEl7y6qHWeO0zIwPAeqjgL2yXsUYTirTZxl8ibCD+JGvKAhzMWteC3GeFJjKS7ErUgCBySPd8+1n/5I1SSJGOoVWKuLxBy0R/OtiMYxtmqsxBvRiY/iRCQ5b3xdxXPsLlU+Rbal6oXYBZa8u6h2nDlOy8BEHBMseY8pGuexl7NN3hPFlyTynbjx9ZBM1PxzXKvmxDtvoQxtfJSRymDJu6zKx20S64PoROIwE8u4kuGtyFiwEIkV7bL2xIW+ABJ9FhotS0veLQU74nCntWAijgmWvMcUjfPYyysd88/ElTi214O21Ay/0rC82dDBR6VawZL34DJ/tACJ+9ZbL1U++fq1L9sUlrzbqHXcsU7LwbFdVpa8FRTDURWw5N2j3J/J11ciyZv6VDQcRwF+N07SBiTsjXgrBpuqjWHJu7FURx/otCK05C0RDGetwItI3n+qEG7FdxFT1Y+NVAtmIiWciobjKDAvLdMmcf9zzY2fHd4I/EHJ3ymbPFO8IaQcEgs5z3dwqr5Tw2kDMBHHBPvkPaZonMdeXkTy/geKVVHi3zyP+Nkpe1Dg90rPDs/S3+nBb18uMjk6RvJ2WmcupuJKZM3U06lkH6eG0wZgIo4JlrzHFI3z2MuLSN7nESo75bkqcKzkvYoEzlUneQfw6wf+KPDUcNoAfKnJmz/STD1VNAJazBqNtEFdFEg1CbbR+ELjJ+KQsOQ9pLrm2xToQYFMPpp+8i40ti3xD+bPP3/4kauIkzfWOT/2gAsrF5fiEMneyS9MxDGhj0/e3+hA6HYrrsRH8VLchVSdG9GJzDX0pwDaO5HXE8/yg9hE41zjNiLxS8WhYMl7KGXNrynQkwIkVxJpSLK73K79WMY3ZZXf3M9X0QqFRq9E/nsIrN83nBzCRBwTDk3euQ7DhR+DhPEYGyrq6LwQ0dqJhv4UQPvLkrut2jzjdcjUQSxufTnEG9iwtiXvoISVpsBIFQgXAuU+hLFc6skefq5+J1Yhl5FLqA0uNJg/9luISzEX+4aTQ8jZxoRDkzeJ4ql0IPQkBuhZhUxG+q9ExkxFQz8KpHKDtrkYo1CDZ7wOK3U4cSqSwIeEJe8h1TXfpkAPCmTywUVC2QROg5qMv9e4Op+596GiMa41ssm6jR1WDHSywUQcEw5N3miXlg7EGdGzKNlDM1eFfkP/CkzkkuRLGYNYoDlvmMoI8VqWOwZqW/IeSFhzawr0pUAmR22SYhjvdmyAMbv6c/W3TQxFhzma0gpOo2EijgmHJm/OMhdX4r04F0kcxGAjlpHKwLcrkPpUNPSrwEzucvGNSAymYiFWvRb5lmTu+75WmYpDw5L30Aqbf1PgQAUyza+6MHa5dXvmrNWP3zJIGJ+J34isSf1S3IVcnffiB5E51OEQcHIKE3FMOCR5ozl6vRVJAqAQlyJ6Uo+B3Yn0bUUnrkRDfwrcyNVWpORZI0Yb0YnoHuKk6vPrw6ncivS9E514LQ6JV3L+yf/nUYcUyHybAqdWINMGuBQomyLMua+YMJPNVdgxTcVcLKJyrvo+JBrwKK5F6kPByTEccg25b41DkjcxcqUVOR8xh4tSH82ZSN8dDUOvClzJW5XuxAH7U81que8ndseAJe9jqGxrmAIHKJBpLpcGZRs4Da6at5a9ra996/JJhLWKfQMP7HeaDxNxTOiavEm+6HZTOgznww7RtgziRx+loV8F+FXEtsJl0Hxd0YfJiQ9UjgRL3kcS2pYxBboqEC6Nthd1mHcfLTxT3UXtvqphrbZ7bLs+e4eJOCZ0Td5OhyAJl88T3gzxVW0VChmr5lWNNVtzBXh+0XVVMSVoXn6jxVDix6+N1jSOBEveRxLaljEFuioQLhTKtninCVxGYe5a9UVbJw3GFxrDOlWfEhtMbzzEaSRMxDHh32kznP83WmyKMzDnXcWche/LKvowvRWrPuVN6fSYqIQB01CxslaBO/UQE/Qvo+4TOeNmIvOYH2NXDHitxPGJ5zWpk7z/n/gfxaTJBBtjCpgCx1WAC5yLgbItwtx7TZyJ27YOGo7faNx7MYnGT1Vfil+LN2IuMu7S12nT1wZOg2G8jponxx9rB8To11rshDMw523FHOKFVlVg3qO4LnXmaqP3G1/eqnwQr307V4lfQ70CJF9ikpWG0MbOc1yF0F+eRzyIwVdiIdImBtRZy4nEpwsSTWJPf9Jlss0xBUyB4RXgQuBFStkF7zSJ+XxyWHRxsGdOov6tuC6Ny9Weiqwd1s1V5zKbiIC+UH827Pnh1A9Zc0wotBnO0nZfK815FEPSZz62WCM1f4ILtViLyz/GUo1rcSsGTZ3q2MFUZJ6hXoGZup7EP4yG/Fz1D2IR2cpV+srxJxap6MRbEUxF/F+KgHm7/D4PqvmRyM6aXefXuDWzKWAK9KVAJke8SCm7gE8LzH8vJl0c7JmDT/yXk8lUtoXoxACnCjYwFbnI2sBpMGTNMaHQZtCg7b4mmrMSH0W+hdiKr0XsdcjUwVpXFQOYy14CSDoz31io3IQOK2sVQFeeyzfiytdvVO4CurqKAYlsxOrC95Vj8K3s+3z7qR8VwXfxUY8ZTAFTYBQKZNoFFwBlV6w08brr5D3zwv649MpYyRBfLpxjKoJbkUuvDZwGw0QcEzgjZ+u6Ly539Gsyn7WexCqQDFLfcaly6+sUaH0rYjfsV4B4zPYPe47Zo8YRlzKmNFjwAAAHEklEQVRSGR4i40p1YgCmIs8M6BKTRPOYX7UuPg2mgClwYgUyrc+LlHJM+Fqb4eLgMnovJmIZXGph35TbaIBTfS7m4kRsAqdBsGotmU+GQisTo6H2NZXvD+JKfCu+FstAQ/YQcKfKKjRU0jcVV6LhcAV47tH0xpezCpc8F3GsHtW+8uPmKjdi6qmiFXjWWJ81DKaAKTBCBTLtiRfpN2Je4lTtU+BCi7InLp97kQusCk+RMVd9FbWXqnOm68gWVzl3XuKj2k7k4hoTuEDRY6h9oQX+5+JWRP8yLmXII+Ot6rG2S98/icZYtbsCG019EHPxtViFXMY06mBOwFSVpXgbDC1LnjWeCZ49gylgCoxQgXBx80ItMz3hfrl4VuJQlwd+y+el7cSxIew1GXBjC/lmnemAa5jr5gpMNJR4xH/c1nz24SMTueD1wB4MpoApMFIFeKFWcaTb7W1bVWfGNjYU2hAX6Rj3NjatbD/9KMCzZsm7Hy3NiylgCpypApa8zzTwJzy2Je8Tim9LmwKmwMtQwJL3y4jjp3QKS96fUrRsr6aAKTBKBSx5V4eF3wt/JV5Wdz9bf6mf1zv6mfvljv5z7bLkfa6Rt3ObAqZAbwq0Td7/Syvz/2IeAz/vTYWPHd3KxO9l3cddzxYSM/2wDg/qoD+tG9CD/fflYwyxYA//teF5LHk3FMqGmQKmgClQp0Db5P3bcvQ3RsK/UHeoHuxT+diIqViHlTqWdZ2yz0V8TMSh8JtyPJZ4/NWGh7Tk3VAoG2YKmAKmQJ0CbZN3nR+zmwJNFbDk3VQpG2cKmAKmQI0ClrxrhDHzYApY8h5MWnNsCpgC56KAJe9zifR4zmnJezyxsJ2YAqbAJ6pA38k7lQ65eCUaTq9Aqi3cimOKhyXv0z8XtgNTwBT4xBXoK3lfSocHEX8L8Z34Rhzyj7Xk3lCjAPFwYjkeNcOPak60Gn+Fz94MpoApYAqYAh0U4ALlIuVCPQT3mrwoOXBqL0s2ax5HgUctk5WW+k7tuv8RSWnooE2eNUveg0pszk0BU+ClK1DogIcmbz7l4WNTEotkjv2iZLfmsAqkco/uKzEGseZ/z0ryPCVYn/2xH4MpYAqYAqZABwW4QLlID73Q5/IxEWNkauCb0nA8BYjDXCzH40a2McSDZ82St0QwmAKmgCnQVYFCE/tI3ny6vhP59M3vui/FhVhOFjPZMnEqAtqpWE40MhkOUIB45CLxgMSjEKvikcoe9D9GPBK/D/ZjMAVMAVPAFOigQLjQuVC74koT+TqW5I0fEsFKXItxslh4263KrUi5EnPxW9HQjwKZ3GzFGzHEY6M6JB7ECxAvYhTikavOmKV4Lw4F9sQ+iqEWML+mgClgCrx0BbhAuUi5ULtgpkkk7vJFnMmGX3gh0n4tBjhVtuJEfBQZZzhcARIzWpKYY8TxSNRxIxbRAKf61reJx5OvD1GwPnuM1x9iHfNpCpgCpsCLVYALlIuUC7ULnCYxnyQcI1MD+4M3zlVOfZ3iUVz79rVKGPCZKl+KvCkwtFOAbzCeKqaEeLz1fbnKSTQOrV/79q3K1NeHKBI55dkohnBuPk0BU8AUOAcFuEC5SLlQ2+JCE5gbEkI8f+H7KMuYycC8qj7GpuJcZIyhuQJXGopm64opd76vSvMw76Zi3hCmxO+lGMK5+TQFTAFT4BwU4ALlwudCbYtME5hbdQmvZOcTYJXfhezMm4l1SNXBGENzBYKuJOoy3srwXkzKHWozHq0vKvowTSM7Y+JP7FFX4yp7qHtuGjuxgaaAKWAKnLMChQ7PRVp1qe/TJfNzb0oDp96+iOy56m98+63KrRjWnKm+FmOkaljyjhXZX0dvNMtKQ6+8PY7T17J948c5ld+JiW8z/7WvX6vMRb6Ov/N8UIm9K1iHfRZdHdg8U8AUMAXOXQEuUC5SLtQu2GrSOpp4oToXfWyjmzW49Cfik8gYQHsjzmhESFVnjqG5AmiPtiHxMvPnIloXNCKgrRMvxVBX9TkeTuVUBLnIGPxORbAUy/6eOxr+SDSONQ/x0XApG2YKmAKmwMtUgAuUi5QLtQuuNGkrkiC+8fWFyjJyGUJSn6q+Ee99SXIoI5WBfRnaKUA8SLQhHtRvKlwsZUP/lTgVN2JoX6oe404N+gMeVanyGfr3lTxrxLbYN9D6TQFTwBQwBaoV4ALlIuVCPQQzTYa70GaNVI7Yl6GbAk3i0dQzifvWD8bvk6+nvmxb8BxY8m6rmo03BUwBUyBSoFCdi/RedBFpc8meCqkWZl+G0yvwQVsgaQM+ha/EuZiKTbDUIFciseXZM5gCpoApYAp0UOALzXlXww7uepmykBcu9sC0F6/mpKsCfPIOmKqyEq/FpiB5Vz1jxNlgCpgCpoApYAqYAqaAKWAKmAKmgClgCpgCR1fg/wMSGBcB2d/qvAAAAABJRU5ErkJggg==) But that's not all :-) You can compose gradients! For instance:<jupyter_code>dfunc_dx = jax.grad(lambda x, A, B: np.sum(jax.jacobian(f2)(x, A, B)))
dfunc_dx(params, A, B).shape<jupyter_output><empty_output><jupyter_text>This is for instance what we used in the tomo challenge to get gradients of the Fisher matrix ;-)  **Something to know about autodiff**: There are two types of autodiff, **forward** and **backward** automatic differentiation. JAX supports both, and you should know the  following:- backward autodiff is good when `f` has small dimensionality of output- forward autodiff is good when `f` has small dimensionality of inputDepending on the relative size of inputs and outputs, you may prefer the dedicated `jax.jacfwd` and `jax.jacrev`Check the [autodiff cookbook](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html) to know more.  Where will JAX suprise you?As mentioned in the previous section, to get all of this goodness, JAX needs tomake some assumptions, which may change a little bit from your habits, and therefore can suprise you at first.  JAX wants pure functionsJAX is based on a purely functional approach, with no side effects. A pure function is a function where the output only depends on the inputs of the function.Let's see what we mean by that by creating a simple jitted function:<jupyter_code>@jax.jit
def my_func(x):
  print("Baguette")
  return 2*np.abs(x)
y = my_func(0) # The first time I execute it I get:
y = my_func(1) #  Second time: I see no print!
y = my_func(2.0) # Third time: The print is back???<jupyter_output>Baguette<jupyter_text>This is because print is a `side effect` and not actually something that depends on the input and affects the output, so it is not included in the compiled graph, it only shows up when JAX traces the function to recompile it.The fact that the third time work again is because JAX detected that the function was being applied on a float now, instead of an int, so it recompiles the function for float inputs.  Your code encounters a tracer instead of an actual value Another example of something suprising:<jupyter_code>@jax.jit
def myfun(x):
  if x > 1:
    y = 0
  else:
    y = x
  return y
myfun(1)<jupyter_output><empty_output><jupyter_text>This doesnt work because when JAX traces the function, it records the execution of the function using an abstract "tracer" instead of a concrete tensor. If your code encouters something that doesn't know how to handle such a tracer, it will trigger an error. To fix this here, you need to use the jax NumPy API:<jupyter_code>@jax.jit
def myfun(x):
  y = np.where(x > 1, 0, x)
  return y

myfun(1)<jupyter_output><empty_output><jupyter_text>Or  use a lower level [jax.lax](https://jax.readthedocs.io/en/latest/jax.lax.html) control flow instruction. This is more advanced...Checkout [this part](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.htmlstructured-control-flow-primitives) of the JAX documentation maybe.$$\begin{array} {r|rr} \hline \\textrm{construct} & \textrm{jit} & \textrm{grad} \\\hline \\textrm{if} & ❌ & ✔ \\\textrm{for} & ✔* & ✔\\\textrm{while} & ✔* & ✔\\\textrm{lax.cond} & ✔ & ✔\\\textrm{lax.while_loop} & ✔ & \textrm{fwd}\\\textrm{lax.fori_loop} & ✔ & \textrm{fwd}\\\textrm{lax.scan} & ✔ & ✔\\\hline\end{array}$$$\ast$ = argument-__value__-independent loop condition - unrolls the loop   In JAX you can't assign values to an array<jupyter_code>myarray = np.arange(10)   # Let's create and array
myarray[5] = 0            # And try to assign a value<jupyter_output><empty_output><jupyter_text>**The reason**: Allowing in-place mutation of variables makes the tracing difficult, so JAX doesn't allow it.**The solution**:- Most of the time, you don't need such assignement, and can rewrite this in other ways- If you really need to, you can use **functional** API tools like [index_update](https://jax.readthedocs.io/en/latest/_autosummary/jax.ops.index_update.htmljax.ops.index_update)<jupyter_code>myarray = np.arange(10)
print("Input     ", myarray)

# Solution 1: We might want to set a value to zero because it's masked
# in such a case, we can simply multiply the input array by a mask
mask = np.where(myarray == 5,0,1) 
myarray1 = mask * myarray
print("Solution 1", myarray1)

# Solution 2: If really you need to update, you can use jax.ops.index_update
myarray2 = jax.ops.index_update(myarray, 5, 0.)
print("Solution 2", myarray2)<jupyter_output>Input      [0 1 2 3 4 5 6 7 8 9]
Solution 1 [0 1 2 3 4 0 6 7 8 9]
Solution 2 [0 1 2 3 4 0 6 7 8 9]<jupyter_text>Autodiff showdownHere are a few cool examples of autodiff to compare against Julia. You can differentiate with respect to objects!This is not possible in TensorFlow for instance. So imagine this:<jupyter_code>def myfunc(cosmo):
  return 2* cosmo["sigma8"]**2 + 1
cosmo = {"sigma8": 0.8, "omega_m":0.3}
jax.grad(myfunc)(cosmo)<jupyter_output><empty_output><jupyter_text>In this case this is a simple dictionary, but you can register your own objectsso that they can be recognized by JAX's autodiff. See [here](https://jax.readthedocs.io/en/latest/pytrees.html)In jax-cosmo, you can take derivatives with respect to `objects` like redshift distributions, and cosmology :-)    Examples of useful implementations End-to-End Limber Integration ExampleOne of the most frequent things we need to do usually is evaluating an integral over a power spectrum.<jupyter_code>def simps(f, a, b, N=128):
  """ Simple Simspon's integration 
  Args:
    f: vectorized function providing the integrand
    a,b: float, interval of integraions
    N: integer, number of points to evaluate the integrand 
  """
  if N % 2 == 1:
      raise ValueError("N must be an even integer.")
  dx = (b - a) / N                # Step between points 
  x = np.linspace(a, b, N + 1)    # Grid for the evaluation of integrand
  y = f(x)                        # Evaluate the integrand on appropriate grid
  return dx / 3 * np.sum(y[0:-1:2] + 4 * y[1::2] + y[2::2], axis=0) # Simpson's rule
pip install --quiet jax-cosmo<jupyter_output><empty_output><jupyter_text>Let's try to use it to implement a limber integration:<jupyter_code>import jax_cosmo as jc
cosmo = jc.Planck15()                  # Create a cosmology with default parameters
nz = jc.redshift.smail_nz(1., 2.,  1.) # Create a Smail redshift distribution
tracer = jc.probes.WeakLensing([nz])   # Define a lensing probe
def angular_cl(cosmo, tracer, ell, integration_fn=simps):

  @jax.vmap
  def limber(ell):
    """ Function that will perform a limber integration at a single ell
    """
    def integrand(a):
      # Step 1: retrieve the associated comoving distance, clip to avoid stupid values
      chi = np.clip(jc.background.radial_comoving_distance(cosmo, a), 1e-2)
      
      # Step 2: get the power spectrum for this combination of chi and a
      pk = jc.power.nonlinear_matter_power(cosmo, k = (ell + 0.5) / chi, a = a) 

      # Step 3: retrieve the kernel
      kernel = tracer.kernel(cosmo, jc.utils.a2z(a), ell)

      # Step 4: computing the integrand
      return  pk * kernel**2 * jc.background.dchioverda(cosmo, a) / chi

    return integration_fn(integrand, jc.utils.z2a(tracer.zmax), 1.0) / jc.constants.c**2

  # Now, evaluate the Limber integral for an array of ell
  return limber(ell)<jupyter_output><empty_output><jupyter_text>And that's it folks<jupyter_code>ell = np.logspace(1,3)
cls = angular_cl(cosmo, tracer, ell)
loglog(ell, cls)
# We can also jit it
jitted_angular_cl = jax.jit(angular_cl)
jitted_angular_cl(cosmo, tracer, ell)<jupyter_output>/usr/local/lib/python3.7/dist-packages/jax/_src/numpy/lax_numpy.py:5792: UserWarning: Explicitly requested dtype <class 'jax._src.numpy.lax_numpy.int64'> requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.
  lax._check_user_dtype_supported(dtype, "astype")<jupyter_text>Solving ODEsSorry running out of time for this, but have a look  at how this is done [here](https://github.com/DifferentiableUniverseInitiative/jax_cosmo/blob/c2f8b59214df70d820bbad53f6a7b1ab2280638a/jax_cosmo/background.pyL471)   ConclusionThis is only the very tip of the iceberg!Here are a few pointers to exciting things in JAX:- [flax](https://github.com/google/flax) : Neural Network library- [morphine](https://github.com/benjaminpope/morphine) : Super cool example of a project which has transitioned from numpy to jax, for the purpose of PSF modeling<jupyter_code><jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.06706382385730211, 'fraction_numerical': 0.08305462653288741, 'mean_word_length': 9.100281491907108, 'pattern_counts': {'":': 2, '<': 95, '<?xml version=': 0, '>': 100, 'https://': 15, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '25', 'dir': 'jupyter-structured-clean-dedup', 'id': '181316', 'license': 'MIT', 'n_tokens_mistral': 16275, 'n_tokens_neox': 15087, 'n_words': 1821, 'path': 'notebooks/DESCJaxTutorial.ipynb', 'repo_name': 'LSSTDESC/jax-perlmutter-tutorials', 'size': '28714'}