Created JAX wikipedia article

[bquast]

Hi colleagues,

I created a Wikipedia article (stub) on JAX: https://en.m.wikipedia.org/wiki/Google_JAX

Right now it just has this content:

Google JAX

Google JAX is a machine learning framework for transforming numerical functions.[1][2][3] The primary function that are transformed are[1]:

1. grad: automatic differentiation
2. jit: compilation
3. vmap: auto-vectorization
4. pmap: SPMD programming

References

1. ^ ^a b JAX: Autograd and XLA, Google, 2022-06-18, archived from the original on 2022-06-18, retrieved 2022-06-18
2. ^ "Using JAX to accelerate our research". www.deepmind.com. Archived from the original on 2022-06-18. Retrieved 2022-06-18.
3. ^ "Why is Google's JAX so popular?". Analytics India Magazine. 2022-04-25. Archived from the original on 2022-06-18. Retrieved 2022-06-18.

This Google-related article is a stub. You can help Wikipedia by expanding it.

Stub icon
This artificial intelligence-related article is a stub. You can help Wikipedia by expanding it.

Infobox

With this infobox (template: infobox software):

JAX	-
Developer(s)	Google
Preview	releasev0.3.13 / 16 May 2022; 33 days ago
Repo	github.com/google/jax
Programing language	Python, C++
Middleware	Numpy
Operating system	Linux, macOS, Windows
Platfor	Python
Size	9.0 Mb
Type	machine learning
License	Apache 2.0

I will continue to expand the article as best I can, but I would like to invite more knowledgable colleagues to contribute content that is missing, and also correct any mistakes I made.

My Wikipedia account is also bquast. Please reach out here or elsewhere if interested.

The idea would be that Google Search will also generate a Google Knowledge Graph box on the side of search results of: https://www.google.com/search?q=google+jax

[bquast]

added the logo and a screenshot of the code sample from README.md

[bquast]

I added a section to demo the grad function (on the logistic function)

grad

The below code demonstrates the grad function's automatic differentiation.

# imports
from jax import grad
import jax.numpy as jnp

# define the logistic function
def logistic(x):  
    return jnp.exp(x)/(jnp.exp(x)+1)

# obtain the gradient function of the logistic function
grad_logistic = grad(logistic)

# evaluate the gradient of the logistic function at x = 1 
print(grad_logistic(1.0))   
# outputsː 0.19661194

[bquast]

I added a section to demo the jit function (on the cube function).

jit

The below code demonstrates the jit function's optimization through fusion.

import jax.numpy as jnp
from jax import jit

def cube(x):
  return x * x * x

x = jnp.ones((10000, 10000))
jit_cube = jit(cube)
cube(x)
jit_cube(x)

The computation time for jit_cube should be noticably shorter than that for cube. Increasing the values on line no. 7, will further exacerbate the difference.

[bquast]

I added a section to demo the pmap function (matmul of PRNGKey values:

pmap

The below code demonstrates the pmap function's parallelization for matrix multiplication.

# import pmap and random from JAX; import JAX NumPy
from jax import pmap, random
import jax.numpy as jnp

# generate 2 random matrices of dimensions 5000 x 6000, one per device
random_keys = random.split(random.PRNGKey(0), 2)
matrices = pmap(lambda key: random.normal(key, (5000, 6000)))(random_keys)

# without data transfer, in parallel, perform a local matrix multiplication on each CPU/GPU
outputs = pmap(lambda x: jnp.dot(x, x.T))(matrices)

# without data transfer, in parallel, obtain the mean for both matrices on each CPU/GPU separately
print(pmap(jnp.mean)(outputs))

The final line should print the valuesː

[1.1566595 1.1805978]

[bquast]

I also added a vmap section and example now (from numpyro). However, I do not understand it well. Perhaps someone can look at it and make it more informative.

vmap

The below code demonstrates the vmap function's vectorization.

def grads(self, inputs):
        in_grad_partial = jax.partial(self._net_grads, self._net_params)
        grad_vmap = jax.vmap(in_grad_partial)
        rich_grads = grad_vmap(inputs)
        flat_grads = np.asarray(self._flatten_batch(rich_grads))
        assert flat_grads.ndim == 2 and flat_grads.shape[0] == inputs.shape[0]
        return flat_grads 

[bquast]

I also added an external links section, to the readthedocs, and the colab quickstart guide.