enigmare/v2-crawler

1{"id":"doc-installation_jax_documentation-eb2133ea","source":"documentation","title":"Installation — JAX documentation","url":"https://docs.jax.dev/en/latest/installation.html","text":"Example:\n```text\npip install -U jax\n```\n\nExample:\n```text\npip install -U \"jax[cuda13]\"\n```\n\nExample:\n```text\npip install -U \"jax[rocm7-local]\"\n```\n\nExample:\n```text\npip install -U \"jax[tpu]\"\n```\n\nExample:\n```text\npip install --upgrade pip\npip install --upgrade jax\n```\n\nExample:\n```text\npip install --upgrade pip\n\n# NVIDIA CUDA 13 installation\n# Note: wheels only available on linux.\npip install --upgrade \"jax[cuda13]\"\n\n# Alternatively, for CUDA 12, use\n# pip install --upgrade \"jax[cuda12]\"\n```\n\nExample:\n```text\npip install --upgrade pip\n\n\n# Installs the wheel compatible with NVIDIA CUDA 13 and cuDNN 9.8 or newer.\n# Note: wheels only available on linux.\npip install --upgrade \"jax[cuda13-local]\"\n\n# Installs the wheel compatible with NVIDIA CUDA 12 and cuDNN 9.8 or newer.\n# Note: wheels only available on linux.\n# pip install --upgrade \"jax[cuda12-local]\"\n```\n\nExample:\n```text\nnvcc --version\n```\n\nExample:\n```text\npip install \"jax[tpu]\"\n```\n\nExample:\n```text\npip install --upgrade \"jax[rocm7-local]\"\n```\n\nExample:\n```text\npython3 -c \"import jax; print(jax.devices())\"\n```\n\nExample:\n```text\ndocker pull rocm/jax:latest\n```\n\nExample:\n```text\nconda install jax -c conda-forge\n```\n\nExample:\n```text\nconda install \"jaxlib=*=*cuda*\" jax -c conda-forge\n```\n\nExample:\n```text\npip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/\n```\n\nExample:\n```text\npip install -U --pre jax jaxlib libtpu requests -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\n```\n\nExample:\n```text\npip install -U --pre jax jaxlib \"jax-cuda13-plugin[with-cuda]\" jax-cuda13-pjrt -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/\n```\n\nExample:\n```text\npip install -U --pre jax jaxlib \"jax-cuda12-plugin[with-cuda]\" jax-cuda12-pjrt -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/\n```\n\nExample:\n```text\n# Install jaxlib on CPU via the wheel archive\npip install \"jax[cpu]==0.3.25\" -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/\n\n# Install the jaxlib 0.3.25 CPU wheel directly\npip install jaxlib==0.3.25 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/\n```\n\nExample:\n```text\npip install jaxlib==0.3.25+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.664Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":20,"totalLines":122,"estimatedTokens":595}}2{"id":"doc-just_in_time_compilation_jax_documentation-ffa3a833","source":"documentation","title":"Just-in-time compilation — JAX documentation","url":"https://docs.jax.dev/en/latest/jit-compilation.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n\nglobal_list = []\n\ndef log2(x):\n  global_list.append(x)\n  ln_x = jnp.log(x)\n  ln_2 = jnp.log(2.0)\n  return ln_x / ln_2\n\nprint(jax.make_jaxpr(log2)(3.0))\n```\n\nExample:\n```text\n{ lambda ; a:f32[]. let\n    b:f32[] = log a\n    c:f32[] = log 2.0:f32[]\n    d:f32[] = div b c\n  in (d,) }\n```\n\nExample:\n```text\ndef log2_with_print(x):\n  print(\"printed x:\", x)\n  ln_x = jnp.log(x)\n  ln_2 = jnp.log(2.0)\n  return ln_x / ln_2\n\nprint(jax.make_jaxpr(log2_with_print)(3.))\n```\n\nExample:\n```text\nprinted x: JitTracer(~float32[])\n{ lambda ; a:f32[]. let\n    b:f32[] = log a\n    c:f32[] = log 2.0:f32[]\n    d:f32[] = div b c\n  in (d,) }\n```\n\nExample:\n```text\ndef log2_if_rank_2(x):\n  if x.ndim == 2:\n    ln_x = jnp.log(x)\n    ln_2 = jnp.log(2.0)\n    return ln_x / ln_2\n  else:\n    return x\n\nprint(jax.make_jaxpr(log2_if_rank_2)(jax.numpy.array([1, 2, 3])))\n```\n\nExample:\n```text\n{ lambda ; a:i32[3]. let  in (a,) }\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\n\ndef selu(x, alpha=1.67, lambda_=1.05):\n  return lambda_ * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)\n\nx = jnp.arange(1000000)\n%timeit selu(x).block_until_ready()\n```\n\nExample:\n```text\n3.71 ms ± 147 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n```\n\nExample:\n```text\nselu_jit = jax.jit(selu)\n\n# Pre-compile the function before timing...\nselu_jit(x).block_until_ready()\n\n%timeit selu_jit(x).block_until_ready()\n```\n\nExample:\n```text\n287 μs ± 2.72 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n```\n\nExample:\n```text\n# Condition on value of x.\n\ndef f(x):\n  if x > 0:\n    return x\n  else:\n    return 2 * x\n\njax.jit(f)(10)  # Raises an error\n```\n\nExample:\n```text\nTracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].\nThe error occurred while tracing the function f at /tmp/ipykernel_2111/2956679937.py:3 for jit. This concrete value was not available in Python because it depends on the value of the argument x.\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError\n```\n\nExample:\n```text\n# While loop conditioned on x and n.\n\ndef g(x, n):\n  i = 0\n  while i < n:\n    i += 1\n  return x + i\n\njax.jit(g)(10, 20)  # Raises an error\n```\n\nExample:\n```text\nTracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].\nThe error occurred while tracing the function g at /tmp/ipykernel_2111/722961019.py:3 for jit. This concrete value was not available in Python because it depends on the value of the argument n.\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError\n```\n\nExample:\n```text\n# While loop conditioned on x and n with a jitted body.\n\n@jax.jit\ndef loop_body(prev_i):\n  return prev_i + 1\n\ndef g_inner_jitted(x, n):\n  i = 0\n  while i < n:\n    i = loop_body(i)\n  return x + i\n\ng_inner_jitted(10, 20)\n```\n\nExample:\n```text\nArray(30, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\nf_jit_correct = jax.jit(f, static_argnums=0)\nprint(f_jit_correct(10))\n```\n\nExample:\n```text\n10\n```\n\nExample:\n```text\ng_jit_correct = jax.jit(g, static_argnames=['n'])\nprint(g_jit_correct(10, 20))\n```\n\nExample:\n```text\n30\n```\n\nExample:\n```text\n@jax.jit(static_argnames=['n'])\ndef g_jit_decorated(x, n):\n  i = 0\n  while i < n:\n    i += 1\n  return x + i\n\nprint(g_jit_decorated(10, 20))\n```\n\nExample:\n```text\nfrom functools import partial\n\ndef unjitted_loop_body(prev_i):\n  return prev_i + 1\n\ndef g_inner_jitted_partial(x, n):\n  i = 0\n  while i < n:\n    # Don't do this! each time the partial returns\n    # a function with different hash\n    i = jax.jit(partial(unjitted_loop_body))(i)\n  return x + i\n\ndef g_inner_jitted_lambda(x, n):\n  i = 0\n  while i < n:\n    # Don't do this!, lambda will also return\n    # a function with a different hash\n    i = jax.jit(lambda x: unjitted_loop_body(x))(i)\n  return x + i\n\ndef g_inner_jitted_normal(x, n):\n  i = 0\n  while i < n:\n    # this is OK, since JAX can find the\n    # cached, compiled function\n    i = jax.jit(unjitted_loop_body)(i)\n  return x + i\n\nprint(\"jit called in a loop with partials:\")\n%timeit g_inner_jitted_partial(10, 20).block_until_ready()\n\nprint(\"jit called in a loop with lambdas:\")\n%timeit g_inner_jitted_lambda(10, 20).block_until_ready()\n\nprint(\"jit called in a loop with caching:\")\n%timeit g_inner_jitted_normal(10, 20).block_until_ready()\n```\n\nExample:\n```text\njit called in a loop with partials:\n308 ms ± 5.03 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\njit called in a loop with lambdas:\n306 ms ± 2.31 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\njit called in a loop with caching:\n1.71 ms ± 1.46 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.665Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":23,"totalLines":244,"estimatedTokens":1175}}3{"id":"doc-jax_the_sharp_bits_jax_documentation-bdcf5ce4","source":"documentation","title":"🔪 JAX - The Sharp Bits 🔪 — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html","text":"Example:\n```text\nimport numpy as np\nfrom jax import jit\nfrom jax import lax\nfrom jax import random\nimport jax\nimport jax.numpy as jnp\n```\n\nExample:\n```text\ndef impure_print_side_effect(x):\n  print(\"Executing function\")  # This is a side-effect\n  return x\n\n# The side-effects appear during the first run\nprint (\"First call: \", jit(impure_print_side_effect)(4.))\n\n# Subsequent runs with parameters of same type and shape may not show the side-effect\n# This is because JAX now invokes a cached compilation of the function\nprint (\"Second call: \", jit(impure_print_side_effect)(5.))\n\n# JAX re-runs the Python function when the type or shape of the argument changes\nprint (\"Third call, different type: \", jit(impure_print_side_effect)(jnp.array([5.])))\n```\n\nExample:\n```text\nExecuting function\nFirst call:  4.0\nSecond call:  5.0\nExecuting function\nThird call, different type:  [5.]\n```\n\nExample:\n```text\ng = 0.\ndef impure_uses_globals(x):\n  return x + g\n\n# JAX captures the value of the global during the first run\nprint (\"First call: \", jit(impure_uses_globals)(4.))\ng = 10.  # Update the global\n\n# Subsequent runs may silently use the cached value of the globals\nprint (\"Second call: \", jit(impure_uses_globals)(5.))\n\n# JAX re-runs the Python function when the type or shape of the argument changes\n# This will end up reading the latest value of the global\nprint (\"Third call, different type: \", jit(impure_uses_globals)(jnp.array([4.])))\n```\n\nExample:\n```text\nFirst call:  4.0\nSecond call:  5.0\nThird call, different type:  [14.]\n```\n\nExample:\n```text\ng = 0.\ndef impure_saves_global(x):\n  global g\n  g = x\n  return x\n\n# JAX runs once the transformed function with special Traced values for arguments\nprint (\"First call: \", jit(impure_saves_global)(4.))\nprint (\"Saved global: \", g)  # Saved global has an internal JAX value\n```\n\nExample:\n```text\nFirst call:  4.0\nSaved global:  JitTracer(~float32[])\n```\n\nExample:\n```text\ndef pure_uses_internal_state(x):\n  state = dict(even=0, odd=0)\n  for i in range(10):\n    state['even' if i % 2 == 0 else 'odd'] += x\n  return state['even'] + state['odd']\n\nprint(jit(pure_uses_internal_state)(5.))\n```\n\nExample:\n```text\n50.0\n```\n\nExample:\n```text\nimport jax.numpy as jnp\nfrom jax import make_jaxpr\n\n# lax.fori_loop\narray = jnp.arange(10)\nprint(lax.fori_loop(0, 10, lambda i,x: x+array[i], 0)) # expected result 45\niterator = iter(range(10))\nprint(lax.fori_loop(0, 10, lambda i,x: x+next(iterator), 0)) # unexpected result 0\n\n# lax.scan\ndef func11(arr, extra):\n    ones = jnp.ones(arr.shape)\n    def body(carry, aelems):\n        ae1, ae2 = aelems\n        return (carry + ae1 * ae2 + extra, carry)\n    return lax.scan(body, 0., (arr, ones))\nmake_jaxpr(func11)(jnp.arange(16), 5.)\n# make_jaxpr(func11)(iter(range(16)), 5.) # throws error\n\n# lax.cond\narray_operand = jnp.array([0.])\nlax.cond(True, lambda x: x+1, lambda x: x-1, array_operand)\niter_operand = iter(range(10))\n# lax.cond(True, lambda x: next(x)+1, lambda x: next(x)-1, iter_operand) # throws error\n```\n\nExample:\n```text\n45\n0\n```\n\nExample:\n```text\nnumpy_array = np.zeros((3,3), dtype=np.float32)\nprint(\"original array:\")\nprint(numpy_array)\n\n# In place, mutating update\nnumpy_array[1, :] = 1.0\nprint(\"updated array:\")\nprint(numpy_array)\n```\n\nExample:\n```text\noriginal array:\n[[0. 0. 0.]\n [0. 0. 0.]\n [0. 0. 0.]]\nupdated array:\n[[0. 0. 0.]\n [1. 1. 1.]\n [0. 0. 0.]]\n```\n\nExample:\n```text\n%xmode Minimal\n```\n\nExample:\n```text\nException reporting mode: Minimal\n```\n\nExample:\n```text\njax_array = jnp.zeros((3,3), dtype=jnp.float32)\n\n# In place update of JAX's array will yield an error!\njax_array[1, :] = 1.0\n```\n\nExample:\n```text\nTypeError: JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html\n```\n\nExample:\n```text\njax_array = jnp.array([10, 20])\njax_array_new = jax_array\njax_array_new += 10\nprint(jax_array_new)  # `jax_array_new` is rebound to a new value [20, 30], but...\nprint(jax_array)      # the original value is unmodified as [10, 20] !\n\nnumpy_array = np.array([10, 20])\nnumpy_array_new = numpy_array\nnumpy_array_new += 10\nprint(numpy_array_new)  # `numpy_array_new is numpy_array`, and it was updated\nprint(numpy_array)      # in-place, so both are [20, 30] !\n```\n\nExample:\n```text\n[20 30]\n[10 20]\n[20 30]\n[20 30]\n```\n\nExample:\n```text\njax_array = jnp.zeros((3,3), dtype=jnp.float32)\nupdated_array = jax_array.at[1, :].set(1.0)\nprint(\"updated array:\\n\", updated_array)\n```\n\nExample:\n```text\nupdated array:\n [[0. 0. 0.]\n [1. 1. 1.]\n [0. 0. 0.]]\n```\n\nExample:\n```text\nprint(\"original array unchanged:\\n\", jax_array)\n```\n\nExample:\n```text\noriginal array unchanged:\n [[0. 0. 0.]\n [0. 0. 0.]\n [0. 0. 0.]]\n```\n\nExample:\n```text\nprint(\"original array:\")\njax_array = jnp.ones((5, 6))\nprint(jax_array)\n\nnew_jax_array = jax_array.at[::2, 3:].add(7.)\nprint(\"new array post-addition:\")\nprint(new_jax_array)\n```\n\nExample:\n```text\noriginal array:\n[[1. 1. 1. 1. 1. 1.]\n [1. 1. 1. 1. 1. 1.]\n [1. 1. 1. 1. 1. 1.]\n [1. 1. 1. 1. 1. 1.]\n [1. 1. 1. 1. 1. 1.]]\nnew array post-addition:\n[[1. 1. 1. 8. 8. 8.]\n [1. 1. 1. 1. 1. 1.]\n [1. 1. 1. 8. 8. 8.]\n [1. 1. 1. 1. 1. 1.]\n [1. 1. 1. 8. 8. 8.]]\n```\n\nExample:\n```text\nimport jax.numpy as jnp\nfrom jax import jit\n\nclass CustomClass:\n  def __init__(self, x: jnp.ndarray, mul: bool):\n    self.x = x\n    self.mul = mul\n\n  @jit  # <---- How to do this correctly?\n  def calc(self, y):\n    if self.mul:\n      return self.x * y\n    return y\n```\n\nExample:\n```text\nc = CustomClass(2, True)\nc.calc(3)\n```\n\nExample:\n```text\nTypeError: Error interpreting argument to <function CustomClass.calc at 0x7562b5c374c0> as an abstract array. The problematic value is of type <class '__main__.CustomClass'> and was passed to the function at path self.\nThis typically means that a jit-wrapped function was called with a non-array argument, and this argument was not marked as static using the static_argnums or static_argnames parameters of jax.jit.\n```\n\nExample:\n```text\nclass CustomClass:\n  def __init__(self, x: jnp.ndarray, mul: bool):\n    self.x = x\n    self.mul = mul\n\n  def calc(self, y):\n    return _calc(self.mul, self.x, y)\n\n@jit(static_argnums=0)\ndef _calc(mul, x, y):\n  if mul:\n    return x * y\n  return y\n```\n\nExample:\n```text\nc = CustomClass(2, True)\nprint(c.calc(3))\n```\n\nExample:\n```text\n6\n```\n\nExample:\n```text\nclass CustomClass:\n  def __init__(self, x: jnp.ndarray, mul: bool):\n    self.x = x\n    self.mul = mul\n\n  # WARNING: this example is broken, as we'll see below. Don't copy & paste!\n  @jit(static_argnums=0)\n  def calc(self, y):\n    if self.mul:\n      return self.x * y\n    return y\n```\n\nExample:\n```text\nc.mul = False\nprint(c.calc(3))  # Should print 3\n```\n\nExample:\n```text\nclass CustomClass:\n  def __init__(self, x: jnp.ndarray, mul: bool):\n    self.x = x\n    self.mul = mul\n\n  @jit(static_argnums=0)\n  def calc(self, y):\n    if self.mul:\n      return self.x * y\n    return y\n\n  def __hash__(self):\n    return hash((self.x, self.mul))\n\n  def __eq__(self, other):\n    return (isinstance(other, CustomClass) and\n            (self.x, self.mul) == (other.x, other.mul))\n```\n\nExample:\n```text\nclass CustomClass:\n  def __init__(self, x: jnp.ndarray, mul: bool):\n    self.x = x\n    self.mul = mul\n\n  @jit\n  def calc(self, y):\n    if self.mul:\n      return self.x * y\n    return y\n\n  def _tree_flatten(self):\n    children = (self.x,)  # arrays / dynamic values\n    aux_data = {'mul': self.mul}  # static values\n    return (children, aux_data)\n\n  @classmethod\n  def _tree_unflatten(cls, aux_data, children):\n    return cls(*children, **aux_data)\n\nfrom jax import tree_util\ntree_util.register_pytree_node(CustomClass,\n                               CustomClass._tree_flatten,\n                               CustomClass._tree_unflatten)\n```\n\nExample:\n```text\nc.mul = False  # mutation is detected\nprint(c.calc(3))\n```\n\nExample:\n```text\n3\n```\n\nExample:\n```text\nc = CustomClass(jnp.array(2), True)  # non-hashable x is supported\nprint(c.calc(3))\n```\n\nExample:\n```text\nnp.arange(10)[11]\n```\n\nExample:\n```text\nIndexError: index 11 is out of bounds for axis 0 with size 10\n```\n\nExample:\n```text\njnp.arange(10)[11]\n```\n\nExample:\n```text\nArray(9, dtype=int32)\n```\n\nExample:\n```text\njnp.arange(10.0).at[11].get()\n```\n\nExample:\n```text\nArray(9., dtype=float32)\n```\n\nExample:\n```text\njnp.arange(10.0).at[11].get(mode='fill', fill_value=jnp.nan)\n```\n\nExample:\n```text\nArray(nan, dtype=float32)\n```\n\nExample:\n```text\nnp.sum([1, 2, 3])\n```\n\nExample:\n```text\nnp.int64(6)\n```\n\nExample:\n```text\njnp.sum([1, 2, 3])\n```\n\nExample:\n```text\nTypeError: sum requires ndarray or scalar arguments, got <class 'list'> at position 0.\n```\n\nExample:\n```text\ndef permissive_sum(x):\n  return jnp.sum(jnp.array(x))\n\nx = list(range(10))\npermissive_sum(x)\n```\n\nExample:\n```text\nArray(45, dtype=int32)\n```\n\nExample:\n```text\nmake_jaxpr(permissive_sum)(x)\n```\n\nExample:\n```text\n{ lambda ; a:i32[] b:i32[] c:i32[] d:i32[] e:i32[] f:i32[] g:i32[] h:i32[] i:i32[]\n    j:i32[]. let\n    k:i32[] = convert_element_type[new_dtype=int32 weak_type=False] a\n    l:i32[1] = broadcast_in_dim k\n    m:i32[] = convert_element_type[new_dtype=int32 weak_type=False] b\n    n:i32[1] = broadcast_in_dim m\n    o:i32[] = convert_element_type[new_dtype=int32 weak_type=False] c\n    p:i32[1] = broadcast_in_dim o\n    q:i32[] = convert_element_type[new_dtype=int32 weak_type=False] d\n    r:i32[1] = broadcast_in_dim q\n    s:i32[] = convert_element_type[new_dtype=int32 weak_type=False] e\n    t:i32[1] = broadcast_in_dim s\n    u:i32[] = convert_element_type[new_dtype=int32 weak_type=False] f\n    v:i32[1] = broadcast_in_dim u\n    w:i32[] = convert_element_type[new_dtype=int32 weak_type=False] g\n    x:i32[1] = broadcast_in_dim w\n    y:i32[] = convert_element_type[new_dtype=int32 weak_type=False] h\n    z:i32[1] = broadcast_in_dim y\n    ba:i32[] = convert_element_type[new_dtype=int32 weak_type=False] i\n    bb:i32[1] = broadcast_in_dim ba\n    bc:i32[] = convert_element_type[new_dtype=int32 weak_type=False] j\n    bd:i32[1] = broadcast_in_dim bc\n    be:i32[10] = concatenate[dimension=0] l n p r t v x z bb bd\n    bf:i32[] = reduce_sum[axes=(0,) out_sharding=None] be\n  in (bf,) }\n```\n\nExample:\n```text\njnp.sum(jnp.array(x))\n```\n\nExample:\n```text\ndef nansum(x):\n  mask = ~jnp.isnan(x)  # boolean mask selecting non-nan values\n  x_without_nans = x[mask]\n  return x_without_nans.sum()\n```\n\nExample:\n```text\nx = jnp.array([1, 2, jnp.nan, 3, 4])\nprint(nansum(x))\n```\n\nExample:\n```text\n10.0\n```\n\nExample:\n```text\njax.jit(nansum)(x)\n```\n\nExample:\n```text\nNonConcreteBooleanIndexError: Array boolean indices must be concrete; got bool[5]\n\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.NonConcreteBooleanIndexError\n```\n\nExample:\n```text\n@jax.jit\ndef nansum_2(x):\n  mask = ~jnp.isnan(x)  # boolean mask selecting non-nan values\n  return jnp.where(mask, x, 0).sum()\n\nprint(nansum_2(x))\n```\n\nExample:\n```text\nx = random.uniform(random.key(0), (1000,), dtype=jnp.float64)\nx.dtype\n```\n\nExample:\n```text\n/tmp/ipykernel_4185/1258726447.py:1: UserWarning: Explicitly requested dtype float64 is not available, and will be truncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.\n  x = random.uniform(random.key(0), (1000,), dtype=jnp.float64)\n```\n\nExample:\n```text\ndtype('float32')\n```\n\nExample:\n```text\n# again, this only works on startup!\nimport jax\njax.config.update(\"jax_enable_x64\", True)\n```\n\nExample:\n```text\nimport jax\njax.config.config_with_absl()\n```\n\nExample:\n```text\nimport jax\nif __name__ == '__main__':\n  # calls jax.config.config_with_absl() *and* runs absl parsing\n  jax.config.parse_flags_with_absl()\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax import random\n\njax.config.update(\"jax_enable_x64\", True)\nx = random.uniform(random.key(0), (1000,), dtype=jnp.float64)\nx.dtype # --> dtype('float64')\n```\n\nExample:\n```text\n>>> np.arange(254.0, 258.0).astype('uint8')\narray([254, 255,   0,   1], dtype=uint8)\n\n>>> jnp.arange(254.0, 258.0).astype('uint8')\nArray([254, 255, 255, 255], dtype=uint8)\n```\n\nExample:\n```text\n>>> import jax.numpy as jnp\n>>> subnormal = jnp.float32(1E-45)\n>>> subnormal  # subnormals are representable\nArray(1.e-45, dtype=float32)\n>>> subnormal + 0  # but are flushed to zero within operations\nArray(0., dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.667Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":70,"totalLines":622,"estimatedTokens":3114}}4{"id":"doc-automatic_vectorization_jax_documentation-f643fa1f","source":"documentation","title":"Automatic vectorization — JAX documentation","url":"https://docs.jax.dev/en/latest/automatic-vectorization.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n\nx = jnp.arange(5)\nw = jnp.array([2., 3., 4.])\n\ndef convolve(x, w):\n  output = []\n  for i in range(1, len(x)-1):\n    output.append(jnp.dot(x[i-1:i+2], w))\n  return jnp.array(output)\n\nconvolve(x, w)\n```\n\nExample:\n```text\nArray([11., 20., 29.], dtype=float32)\n```\n\nExample:\n```text\nxs = jnp.stack([x, x])\nws = jnp.stack([w, w])\n```\n\nExample:\n```text\ndef manually_batched_convolve(xs, ws):\n  output = []\n  for i in range(xs.shape[0]):\n    output.append(convolve(xs[i], ws[i]))\n  return jnp.stack(output)\n\nmanually_batched_convolve(xs, ws)\n```\n\nExample:\n```text\nArray([[11., 20., 29.],\n       [11., 20., 29.]], dtype=float32)\n```\n\nExample:\n```text\ndef manually_vectorized_convolve(xs, ws):\n  output = []\n  for i in range(1, xs.shape[-1] -1):\n    output.append(jnp.sum(xs[:, i-1:i+2] * ws, axis=1))\n  return jnp.stack(output, axis=1)\n\nmanually_vectorized_convolve(xs, ws)\n```\n\nExample:\n```text\nauto_batch_convolve = jax.vmap(convolve)\n\nauto_batch_convolve(xs, ws)\n```\n\nExample:\n```text\nauto_batch_convolve_v2 = jax.vmap(convolve, in_axes=1, out_axes=1)\n\nxst = jnp.transpose(xs)\nwst = jnp.transpose(ws)\n\nauto_batch_convolve_v2(xst, wst)\n```\n\nExample:\n```text\nArray([[11., 11.],\n       [20., 20.],\n       [29., 29.]], dtype=float32)\n```\n\nExample:\n```text\nbatch_convolve_v3 = jax.vmap(convolve, in_axes=[0, None])\n\nbatch_convolve_v3(xs, w)\n```\n\nExample:\n```text\njitted_batch_convolve = jax.jit(auto_batch_convolve)\n\njitted_batch_convolve(xs, ws)\n```\n\nExample:\n```text\ndef pairwise(f, xs):\n  return jax.vmap(lambda x: jax.vmap(lambda y: f(x, y))(xs))(xs)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.668Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":12,"totalLines":101,"estimatedTokens":404}}5{"id":"doc-stateful_computations_jax_documentation-fdf36070","source":"documentation","title":"Stateful computations — JAX documentation","url":"https://docs.jax.dev/en/latest/stateful-computations.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n\nclass Counter:\n  \"\"\"A simple counter.\"\"\"\n\n  def __init__(self):\n    self.n = 0\n\n  def count(self) -> int:\n    \"\"\"Increments the counter and returns the new value.\"\"\"\n    self.n += 1\n    return self.n\n\n  def reset(self):\n    \"\"\"Resets the counter to zero.\"\"\"\n    self.n = 0\n\n\ncounter = Counter()\n\nfor _ in range(3):\n  print(counter.count())\n```\n\nExample:\n```text\n1\n2\n3\n```\n\nExample:\n```text\ncounter.reset()\nfast_count = jax.jit(counter.count)\n\nfor _ in range(3):\n  print(fast_count())\n```\n\nExample:\n```text\n1\n1\n1\n```\n\nExample:\n```text\nself.n += 1\n```\n\nExample:\n```text\nCounterState = int\n\nclass CounterV2:\n\n  def count(self, n: CounterState) -> tuple[int, CounterState]:\n    # You could just return n+1, but here we separate its role as \n    # the output and as the counter state for didactic purposes.\n    return n+1, n+1\n\n  def reset(self) -> CounterState:\n    return 0\n\ncounter = CounterV2()\nstate = counter.reset()\n\nfor _ in range(3):\n  value, state = counter.count(state)\n  print(value)\n```\n\nExample:\n```text\nstate = counter.reset()\nfast_count = jax.jit(counter.count)\n\nfor _ in range(3):\n  value, state = fast_count(state)\n  print(value)\n```\n\nExample:\n```text\nclass StatefulClass\n\n  state: State\n\n  def stateful_method(*args, **kwargs) -> Output:\n```\n\nExample:\n```text\nclass StatelessClass\n\n  def stateless_method(state: State, *args, **kwargs) -> (Output, State):\n```\n\nExample:\n```text\nfrom typing import NamedTuple\n\nclass Params(NamedTuple):\n  weight: jnp.ndarray\n  bias: jnp.ndarray\n\n\ndef init(rng) -> Params:\n  \"\"\"Returns the initial model params.\"\"\"\n  weights_key, bias_key = jax.random.split(rng)\n  weight = jax.random.normal(weights_key, ())\n  bias = jax.random.normal(bias_key, ())\n  return Params(weight, bias)\n\n\ndef loss(params: Params, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:\n  \"\"\"Computes the least squares error of the model's predictions on x against y.\"\"\"\n  pred = params.weight * x + params.bias\n  return jnp.mean((pred - y) ** 2)\n\n\nLEARNING_RATE = 0.005\n\n@jax.jit\ndef update(params: Params, x: jnp.ndarray, y: jnp.ndarray) -> Params:\n  \"\"\"Performs one SGD update step on params using the given data.\"\"\"\n  grad = jax.grad(loss)(params, x, y)\n\n  # If we were using Adam or another stateful optimizer,\n  # we would also do something like\n  #\n  #   updates, new_optimizer_state = optimizer(grad, optimizer_state)\n  # \n  # and then use `updates` instead of `grad` to actually update the params.\n  # (And we'd include `new_optimizer_state` in the output, naturally.)\n\n  new_params = jax.tree.map(\n      lambda param, g: param - g * LEARNING_RATE, params, grad)\n\n  return new_params\n```\n\nExample:\n```text\nimport matplotlib.pyplot as plt\n\nrng = jax.random.key(42)\n\n# Generate true data from y = w*x + b + noise\ntrue_w, true_b = 2, -1\nx_rng, noise_rng = jax.random.split(rng)\nxs = jax.random.normal(x_rng, (128, 1))\nnoise = jax.random.normal(noise_rng, (128, 1)) * 0.5\nys = xs * true_w + true_b + noise\n\n# Fit regression\nparams = init(rng)\nfor _ in range(1000):\n  params = update(params, xs, ys)\n\nplt.scatter(xs, ys)\nplt.plot(xs, params.weight * xs + params.bias, c='red', label='Model Prediction')\nplt.legend();\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.668Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":11,"totalLines":171,"estimatedTokens":801}}6{"id":"doc-quickstart_how_to_think_in_jax_jax_documentation-e25bdbef","source":"documentation","title":"Quickstart: How to think in JAX — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html","text":"Example:\n```text\npip install jax\n```\n\nExample:\n```text\npip install -U \"jax[cuda13]\"\n```\n\nExample:\n```text\nimport jax.numpy as jnp\n```\n\nExample:\n```text\nimport matplotlib.pyplot as plt\n\nx_jnp = jnp.linspace(0, 10, 1000)\ny_jnp = 2 * jnp.sin(x_jnp) * jnp.cos(x_jnp)\nplt.plot(x_jnp, y_jnp);\n```\n\nExample:\n```text\nimport numpy as np\nimport jax.numpy as jnp\n\nx_np = np.linspace(0, 10, 1000)\nx_jnp = jnp.linspace(0, 10, 1000)\n```\n\nExample:\n```text\ntype(x_np)\n```\n\nExample:\n```text\nnumpy.ndarray\n```\n\nExample:\n```text\ntype(x_jnp)\n```\n\nExample:\n```text\njaxlib._jax.ArrayImpl\n```\n\nExample:\n```text\n# NumPy: mutable arrays\nx = np.arange(10)\nx[0] = 10\nprint(x)\n```\n\nExample:\n```text\n[10  1  2  3  4  5  6  7  8  9]\n```\n\nExample:\n```text\n%xmode minimal\n```\n\nExample:\n```text\nException reporting mode: Minimal\n```\n\nExample:\n```text\n# JAX: immutable arrays\nx = jnp.arange(10)\nx[0] = 10\n```\n\nExample:\n```text\nTypeError: JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html\n```\n\nExample:\n```text\ny = x.at[0].set(10)\nprint(x)\nprint(y)\n```\n\nExample:\n```text\n[0 1 2 3 4 5 6 7 8 9]\n[10  1  2  3  4  5  6  7  8  9]\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\n\nx = jnp.arange(5)\nisinstance(x, jax.Array)\n```\n\nExample:\n```text\nTrue\n```\n\nExample:\n```text\nx.devices()\n```\n\nExample:\n```text\n{CpuDevice(id=0)}\n```\n\nExample:\n```text\nx.sharding\n```\n\nExample:\n```text\nSingleDeviceSharding(device=CpuDevice(id=0), memory_kind=device)\n```\n\nExample:\n```text\nimport jax.numpy as jnp\n\ndef norm(X):\n  X = X - X.mean(0)\n  return X / X.std(0)\n```\n\nExample:\n```text\nfrom jax import jit\nnorm_compiled = jit(norm)\n```\n\nExample:\n```text\nnp.random.seed(1701)\nX = jnp.array(np.random.rand(10000, 10))\nnp.allclose(norm(X), norm_compiled(X), atol=1E-6)\n```\n\nExample:\n```text\n%timeit norm(X).block_until_ready()\n%timeit norm_compiled(X).block_until_ready()\n```\n\nExample:\n```text\n227 μs ± 840 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n173 μs ± 1.48 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n```\n\nExample:\n```text\ndef get_negatives(x):\n  return x[x < 0]\n\nx = jnp.array(np.random.randn(10))\nget_negatives(x)\n```\n\nExample:\n```text\nArray([-0.10570311, -0.59403396, -0.8680282 , -0.23489487], dtype=float32)\n```\n\nExample:\n```text\njit(get_negatives)(x)\n```\n\nExample:\n```text\nNonConcreteBooleanIndexError: Array boolean indices must be concrete; got bool[10]\n\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.NonConcreteBooleanIndexError\n```\n\nExample:\n```text\nfrom jax import grad\n\ndef sum_logistic(x):\n  return jnp.sum(1.0 / (1.0 + jnp.exp(-x)))\n\nx_small = jnp.arange(3.)\nderivative_fn = grad(sum_logistic)\nprint(derivative_fn(x_small))\n```\n\nExample:\n```text\n[0.25       0.19661197 0.10499357]\n```\n\nExample:\n```text\ndef first_finite_differences(f, x, eps=1E-3):\n  return jnp.array([(f(x + eps * v) - f(x - eps * v)) / (2 * eps)\n                   for v in jnp.eye(len(x))])\n\nprint(first_finite_differences(sum_logistic, x_small))\n```\n\nExample:\n```text\n[0.24998187 0.1964569  0.10502338]\n```\n\nExample:\n```text\nprint(grad(jit(grad(jit(grad(sum_logistic)))))(1.0))\n```\n\nExample:\n```text\n-0.0353256\n```\n\nExample:\n```text\nfrom jax import jacobian\nprint(jacobian(jnp.exp)(x_small))\n```\n\nExample:\n```text\n[[1.        0.        0.       ]\n [0.        2.7182817 0.       ]\n [0.        0.        7.389056 ]]\n```\n\nExample:\n```text\nfrom jax import jacfwd, jacrev\ndef hessian(fun):\n  return jit(jacfwd(jacrev(fun)))\nprint(hessian(sum_logistic)(x_small))\n```\n\nExample:\n```text\n[[-0.         -0.         -0.        ]\n [-0.         -0.09085776 -0.        ]\n [-0.         -0.         -0.07996249]]\n```\n\nExample:\n```text\nfrom jax import random\n\nkey = random.key(1701)\nkey1, key2 = random.split(key)\nmat = random.normal(key1, (150, 100))\nbatched_x = random.normal(key2, (10, 100))\n\ndef apply_matrix(x):\n  return jnp.dot(mat, x)\n```\n\nExample:\n```text\ndef naively_batched_apply_matrix(v_batched):\n  return jnp.stack([apply_matrix(v) for v in v_batched])\n\nprint('Naively batched')\n%timeit naively_batched_apply_matrix(batched_x).block_until_ready()\n```\n\nExample:\n```text\nNaively batched\n197 μs ± 1.62 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n```\n\nExample:\n```text\nimport numpy as np\n\n@jit\ndef batched_apply_matrix(batched_x):\n  return jnp.dot(batched_x, mat.T)\n\nnp.testing.assert_allclose(naively_batched_apply_matrix(batched_x),\n                           batched_apply_matrix(batched_x), atol=1E-4, rtol=1E-4)\nprint('Manually batched')\n%timeit batched_apply_matrix(batched_x).block_until_ready()\n```\n\nExample:\n```text\nManually batched\n14.3 μs ± 193 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n```\n\nExample:\n```text\nfrom jax import vmap\n\n@jit\ndef vmap_batched_apply_matrix(batched_x):\n  return vmap(apply_matrix)(batched_x)\n\nnp.testing.assert_allclose(naively_batched_apply_matrix(batched_x),\n                           vmap_batched_apply_matrix(batched_x), atol=1E-4, rtol=1E-4)\nprint('Auto-vectorized with vmap')\n%timeit vmap_batched_apply_matrix(batched_x).block_until_ready()\n```\n\nExample:\n```text\nAuto-vectorized with vmap\n15.2 μs ± 131 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n```\n\nExample:\n```text\nfrom jax import random\n\nkey = random.key(43)\nprint(key)\n```\n\nExample:\n```text\nArray((), dtype=key<fry>) overlaying:\n[ 0 43]\n```\n\nExample:\n```text\nprint(random.normal(key))\nprint(random.normal(key))\n```\n\nExample:\n```text\n0.07520543\n0.07520543\n```\n\nExample:\n```text\nfor i in range(3):\n  new_key, subkey = random.split(key)\n  del key  # The old key is consumed by split() -- we must never use it again.\n\n  val = random.normal(subkey)\n  del subkey  # The subkey is consumed by normal().\n\n  print(f\"draw {i}: {val}\")\n  key = new_key  # new_key is safe to use in the next iteration.\n```\n\nExample:\n```text\ndraw 0: -1.9133632183074951\ndraw 1: -1.4749839305877686\ndraw 2: -0.36703771352767944\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\n\n@jax.jit\ndef f(x):\n  print(\"print(x) ->\", x)\n  y = jnp.sin(x)\n  print(\"print(y) ->\", y)\n  return y\n\nresult = f(2.)\n```\n\nExample:\n```text\nprint(x) -> JitTracer(~float32[])\nprint(y) -> JitTracer(~float32[])\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  jax.debug.print(\"jax.debug.print(x) -> {x}\", x=x)\n  y = jnp.sin(x)\n  jax.debug.print(\"jax.debug.print(y) -> {y}\", y=y)\n  return y\n\nresult = f(2.)\n```\n\nExample:\n```text\njax.debug.print(x) -> 2.0\njax.debug.print(y) -> 0.9092974066734314\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.670Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":59,"totalLines":418,"estimatedTokens":1636}}7{"id":"doc-tracing_jax_documentation-a801668b","source":"documentation","title":"Tracing — JAX documentation","url":"https://docs.jax.dev/en/latest/tracing.html","text":"Example:\n```text\nfrom jax import jit\nimport jax.numpy as jnp\nimport numpy as np\n\n@jit\ndef f(x, y):\n  print(\"Running f():\")\n  print(f\"  x = {x}\")\n  print(f\"  y = {y}\")\n  result = jnp.dot(x + 1, y + 1)\n  print(f\"  result = {result}\")\n  return result\n\nx = np.random.randn(3, 4)\ny = np.random.randn(4)\nf(x, y)\n```\n\nExample:\n```text\nRunning f():\n  x = JitTracer(float32[3,4])\n  y = JitTracer(float32[4])\n  result = JitTracer(float32[3])\n```\n\nExample:\n```text\nArray([ 0.05875492, 17.501205  , 16.9168    ], dtype=float32)\n```\n\nExample:\n```text\nx2 = np.random.randn(3, 4)\ny2 = np.random.randn(4)\nf(x2, y2)\n```\n\nExample:\n```text\nArray([3.1219645, 2.36379  , 1.5909786], dtype=float32)\n```\n\nExample:\n```text\nfrom jax import make_jaxpr\n\ndef f(x, y):\n  return jnp.dot(x + 1, y + 1)\n\nmake_jaxpr(f)(x, y)\n```\n\nExample:\n```text\n{ lambda ; a:f32[3,4] b:f32[4]. let\n    c:f32[3,4] = add a 1.0:f32[]\n    d:f32[4] = add b 1.0:f32[]\n    e:f32[3] = dot_general[\n      dimension_numbers=(([1], [0]), ([], []))\n      preferred_element_type=float32\n    ] c d\n  in (e,) }\n```\n\nExample:\n```text\n@jit\ndef f(x, neg):\n  return -x if neg else x\n\nf(1, True)\n```\n\nExample:\n```text\n---------------------------------------------------------------------------\nTracerBoolConversionError                 Traceback (most recent call last)\nCell In[4], line 5\n      1 @jit\n      2 def f(x, neg):\n      3   return -x if neg else x\n      4 \n----> 5 f(1, True)\n\n    [... skipping hidden 5 frame]\n\nCell In[4], line 3, in f(x, neg)\n      1 @jit\n      2 def f(x, neg):\n----> 3   return -x if neg else x\n\n    [... skipping hidden 1 frame]\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:2022, in concretization_function_error.<locals>.error(self, arg)\n   2021 def error(self, arg):\n-> 2022   raise TracerBoolConversionError(arg)\n\nTracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].\nThe error occurred while tracing the function f at /tmp/ipykernel_5796/2422663986.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument neg.\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError\n```\n\nExample:\n```text\nfrom functools import partial\n\n@jit(static_argnums=(1,))\ndef f(x, neg):\n  return -x if neg else x\n\nf(1, True)\n```\n\nExample:\n```text\nArray(-1, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\nf(1, False)\n```\n\nExample:\n```text\nArray(1, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\nimport jax.numpy as jnp\nfrom jax import jit\n\n@jit\ndef f(x):\n  return x.reshape(jnp.array(x.shape).prod())\n\nx = jnp.ones((2, 3))\nf(x)\n```\n\nExample:\n```text\n---------------------------------------------------------------------------\nTypeError                                 Traceback (most recent call last)\nCell In[7], line 9\n      5 def f(x):\n      6   return x.reshape(jnp.array(x.shape).prod())\n      7 \n      8 x = jnp.ones((2, 3))\n----> 9 f(x)\n\n    [... skipping hidden 5 frame]\n\nCell In[7], line 6, in f(x)\n      4 @jit\n      5 def f(x):\n----> 6   return x.reshape(jnp.array(x.shape).prod())\n\n    [... skipping hidden 2 frame]\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:522, in _compute_newshape(arr, newshape)\n    520 else:\n    521   newshape: Sequence[DimSize]  # pyrefly: ignore[redefinition]\n--> 522 newshape = core.canonicalize_shape(newshape)\n    523 neg1s = [i for i, d in enumerate(newshape) if type(d) is int and d == -1]\n    524 if len(neg1s) > 1:\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:2125, in canonicalize_shape(shape, context)\n   2123 except TypeError:\n   2124   pass\n-> 2125 raise _invalid_shape_error(shape, context)\n\nTypeError: Shapes must be 1D sequences of concrete values of integer type, got [JitTracer(int32[])].\nIf using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.\nThe error occurred while tracing the function f at /tmp/ipykernel_5796/1983583872.py:4 for jit. This value became a tracer due to JAX operations on these lines:\n\n  operation a:i32[] = reduce_prod[axes=(0,)] b\n    from line /tmp/ipykernel_5796/1983583872.py:6:19 (f)\n```\n\nExample:\n```text\n@jit\ndef f(x):\n  print(f\"x = {x}\")\n  print(f\"x.shape = {x.shape}\")\n  print(f\"jnp.array(x.shape).prod() = {jnp.array(x.shape).prod()}\")\n  # comment this out to avoid the error:\n  # return x.reshape(jnp.array(x.shape).prod())\n\nf(x)\n```\n\nExample:\n```text\nx = JitTracer(float32[2,3])\nx.shape = (2, 3)\njnp.array(x.shape).prod() = JitTracer(int32[])\n```\n\nExample:\n```text\nfrom jax import jit\nimport jax.numpy as jnp\nimport numpy as np\n\n@jit\ndef f(x):\n  return x.reshape((np.prod(x.shape),))\n\nf(x)\n```\n\nExample:\n```text\nArray([1., 1., 1., 1., 1., 1.], dtype=float32)\n```\n\nExample:\n```text\ndef divide(x, y):\n  return x / y if y >= 1. else 0.\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.670Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":20,"totalLines":229,"estimatedTokens":1235}}8{"id":"doc-pytrees_jax_documentation-a33da68f","source":"documentation","title":"Pytrees — JAX documentation","url":"https://docs.jax.dev/en/latest/pytrees.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n\nexample_trees = [\n    [1, 'a', object()],\n    (1, (2, 3), ()),\n    [1, {'k1': 2, 'k2': (3, 4)}, 5],\n    {'a': 2, 'b': (2, 3)},\n    jnp.array([1, 2, 3]),\n]\n\n# Print how many leaves the pytrees have.\nfor pytree in example_trees:\n  # This `jax.tree.leaves()` method extracts the flattened leaves from the pytrees.\n  leaves = jax.tree.leaves(pytree)\n  print(f\"{repr(pytree):<45} has {len(leaves)} leaves: {leaves}\")\n```\n\nExample:\n```text\n[1, 'a', <object object at 0x73b7d4ff55e0>]   has 3 leaves: [1, 'a', <object object at 0x73b7d4ff55e0>]\n(1, (2, 3), ())                               has 3 leaves: [1, 2, 3]\n[1, {'k1': 2, 'k2': (3, 4)}, 5]               has 5 leaves: [1, 2, 3, 4, 5]\n{'a': 2, 'b': (2, 3)}                         has 3 leaves: [2, 2, 3]\nArray([1, 2, 3], dtype=int32)                 has 1 leaves: [Array([1, 2, 3], dtype=int32)]\n```\n\nExample:\n```text\nlist_of_lists = [\n    [1, 2, 3],\n    [1, 2],\n    [1, 2, 3, 4]\n]\n\njax.tree.map(lambda x: x*2, list_of_lists)\n```\n\nExample:\n```text\n[[2, 4, 6], [2, 4], [2, 4, 6, 8]]\n```\n\nExample:\n```text\nanother_list_of_lists = list_of_lists\njax.tree.map(lambda x, y: x+y, list_of_lists, another_list_of_lists)\n```\n\nExample:\n```text\nimport numpy as np\n\ndef init_mlp_params(layer_widths):\n  params = []\n  for n_in, n_out in zip(layer_widths[:-1], layer_widths[1:]):\n    params.append(\n        dict(weights=np.random.normal(size=(n_in, n_out)) * np.sqrt(2/n_in),\n             biases=np.ones(shape=(n_out,))\n            )\n    )\n  return params\n\nparams = init_mlp_params([1, 128, 128, 1])\n```\n\nExample:\n```text\njax.tree.map(lambda x: x.shape, params)\n```\n\nExample:\n```text\n[{'biases': (128,), 'weights': (1, 128)},\n {'biases': (128,), 'weights': (128, 128)},\n {'biases': (1,), 'weights': (128, 1)}]\n```\n\nExample:\n```text\n# Define the forward pass.\ndef forward(params, x):\n  *hidden, last = params\n  for layer in hidden:\n    x = jax.nn.relu(x @ layer['weights'] + layer['biases'])\n  return x @ last['weights'] + last['biases']\n\n# Define the loss function.\ndef loss_fn(params, x, y):\n  return jnp.mean((forward(params, x) - y) ** 2)\n\n# Set the learning rate.\nLEARNING_RATE = 0.0001\n\n# Using the stochastic gradient descent, define the parameter update function.\n# Apply `@jax.jit` for JIT compilation (speed).\n@jax.jit\ndef update(params, x, y):\n  # Calculate the gradients with `jax.grad`.\n  grads = jax.grad(loss_fn)(params, x, y)\n  # Note that `grads` is a pytree with the same structure as `params`.\n  # `jax.grad` is one of many JAX functions that has\n  # built-in support for pytrees.\n  # This is useful - you can apply the SGD update using JAX pytree utilities.\n  return jax.tree.map(\n      lambda p, g: p - LEARNING_RATE * g, params, grads\n  )\n```\n\nExample:\n```text\nfrom jax.tree_util import tree_structure\nprint(tree_structure(object))\n```\n\nExample:\n```text\nPyTreeDef(*)\n```\n\nExample:\n```text\nvmap(f, in_axes=(a1, {\"k1\": a2, \"k2\": a3}))\n```\n\nExample:\n```text\nvmap(f, in_axes=(None, {\"k1\": None, \"k2\": 0}))\n```\n\nExample:\n```text\nvmap(f, in_axes=(None, 0))  # equivalent to (None, {\"k1\": 0, \"k2\": 0})\n```\n\nExample:\n```text\nvmap(f, in_axes=0)  # equivalent to (0, {\"k1\": 0, \"k2\": 0})\n```\n\nExample:\n```text\nimport collections\n\nATuple = collections.namedtuple(\"ATuple\", ('name'))\n\ntree = [1, {'k1': 2, 'k2': (3, 4)}, ATuple('foo')]\nflattened, _ = jax.tree_util.tree_flatten_with_path(tree)\n\nfor key_path, value in flattened:\n  print(f'Value of tree{jax.tree_util.keystr(key_path)}: {value}')\n```\n\nExample:\n```text\nValue of tree[0]: 1\nValue of tree[1]['k1']: 2\nValue of tree[1]['k2'][0]: 3\nValue of tree[1]['k2'][1]: 4\nValue of tree[2].name: foo\n```\n\nExample:\n```text\nfor key_path, _ in flattened:\n  print(f'Key path of tree{jax.tree_util.keystr(key_path)}: {repr(key_path)}')\n```\n\nExample:\n```text\nKey path of tree[0]: (SequenceKey(idx=0),)\nKey path of tree[1]['k1']: (SequenceKey(idx=1), DictKey(key='k1'))\nKey path of tree[1]['k2'][0]: (SequenceKey(idx=1), DictKey(key='k2'), SequenceKey(idx=0))\nKey path of tree[1]['k2'][1]: (SequenceKey(idx=1), DictKey(key='k2'), SequenceKey(idx=1))\nKey path of tree[2].name: (SequenceKey(idx=2), GetAttrKey(name='name'))\n```\n\nExample:\n```text\na_tree = [jnp.zeros((2, 3)), jnp.zeros((3, 4))]\n\n# Try to make another pytree with ones instead of zeros.\nshapes = jax.tree.map(lambda x: x.shape, a_tree)\njax.tree.map(jnp.ones, shapes)\n```\n\nExample:\n```text\n[(Array([1., 1.], dtype=float32), Array([1., 1., 1.], dtype=float32)),\n (Array([1., 1., 1.], dtype=float32), Array([1., 1., 1., 1.], dtype=float32))]\n```\n\nExample:\n```text\njax.tree.leaves([None, None, None])\n```\n\nExample:\n```text\njax.tree.leaves([None, None, None], is_leaf=lambda x: x is None)\n```\n\nExample:\n```text\n[None, None, None]\n```\n\nExample:\n```text\njax.tree.map(lambda x: x + 1, {1: 7, \"y\": 42})\n```\n\nExample:\n```text\nValueError: Comparator raised exception while sorting pytree dictionary keys.\n```\n\nExample:\n```text\ndef tree_transpose(list_of_trees):\n  \"\"\"\n  Converts a list of trees of identical structure into a single tree of lists.\n  \"\"\"\n  return jax.tree.map(lambda *xs: list(xs), *list_of_trees)\n\n# Convert a dataset from row-major to column-major.\nepisode_steps = [dict(t=1, obs=3), dict(t=2, obs=4)]\ntree_transpose(episode_steps)\n```\n\nExample:\n```text\n{'obs': [3, 4], 't': [1, 2]}\n```\n\nExample:\n```text\njax.tree.transpose(\n  outer_treedef = jax.tree.structure([0 for e in episode_steps]),\n  inner_treedef = jax.tree.structure(episode_steps[0]),\n  pytree_to_transpose = episode_steps\n)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.672Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":29,"totalLines":247,"estimatedTokens":1378}}9{"id":"doc-pseudorandom_numbers_jax_documentation-2e5843a1","source":"documentation","title":"Pseudorandom numbers — JAX documentation","url":"https://docs.jax.dev/en/latest/random-numbers.html","text":"Example:\n```text\nimport numpy as np\nnp.random.seed(0)\n```\n\nExample:\n```text\nprint(np.random.random())\nprint(np.random.random())\nprint(np.random.random())\n```\n\nExample:\n```text\n0.5488135039273248\n0.7151893663724195\n0.6027633760716439\n```\n\nExample:\n```text\ndef print_truncated_random_state():\n  \"\"\"To avoid spamming the outputs, print only part of the state.\"\"\"\n  full_random_state = np.random.get_state()\n  print(str(full_random_state)[:460], '...')\n\nprint_truncated_random_state()\n```\n\nExample:\n```text\n('MT19937', array([2443250962, 1093594115, 1878467924, 2709361018, 1101979660,\n       3904844661,  676747479, 2085143622, 1056793272, 3812477442,\n       2168787041,  275552121, 2696932952, 3432054210, 1657102335,\n       3518946594,  962584079, 1051271004, 3806145045, 1414436097,\n       2032348584, 1661738718, 1116708477, 2562755208, 3176189976,\n        696824676, 2399811678, 3992505346,  569184356, 2626558620,\n        136797809, 4273176064,  296167901, 343 ...\n```\n\nExample:\n```text\nnp.random.seed(0)\nprint_truncated_random_state()\n```\n\nExample:\n```text\n('MT19937', array([         0,          1, 1812433255, 1900727105, 1208447044,\n       2481403966, 4042607538,  337614300, 3232553940, 1018809052,\n       3202401494, 1775180719, 3192392114,  594215549,  184016991,\n        829906058,  610491522, 3879932251, 3139825610,  297902587,\n       4075895579, 2943625357, 3530655617, 1423771745, 2135928312,\n       2891506774, 1066338622,  135451537,  933040465, 2759011858,\n       2273819758, 3545703099, 2516396728, 127 ...\n```\n\nExample:\n```text\n_ = np.random.uniform()\nprint_truncated_random_state()\n```\n\nExample:\n```text\nnp.random.seed(0)\nprint(np.random.uniform(size=3))\n```\n\nExample:\n```text\n[0.5488135  0.71518937 0.60276338]\n```\n\nExample:\n```text\nnp.random.seed(0)\nprint(\"individually:\", np.stack([np.random.uniform() for _ in range(3)]))\n\nnp.random.seed(0)\nprint(\"all at once: \", np.random.uniform(size=3))\n```\n\nExample:\n```text\nindividually: [0.5488135  0.71518937 0.60276338]\nall at once:  [0.5488135  0.71518937 0.60276338]\n```\n\nExample:\n```text\nimport numpy as np\n\nnp.random.seed(0)\n\ndef bar(): return np.random.uniform()\ndef baz(): return np.random.uniform()\n\ndef foo(): return bar() + 2 * baz()\n\nprint(foo())\n```\n\nExample:\n```text\n1.9791922366721637\n```\n\nExample:\n```text\nfrom jax import random\n\nkey = random.key(42)\nprint(key)\n```\n\nExample:\n```text\nArray((), dtype=key<fry>) overlaying:\n[ 0 42]\n```\n\nExample:\n```text\nprint(random.normal(key))\nprint(random.normal(key))\n```\n\nExample:\n```text\n-0.028304616\n-0.028304616\n```\n\nExample:\n```text\nfor i in range(3):\n  new_key, subkey = random.split(key)\n  del key  # The old key is consumed by split() -- we must never use it again.\n\n  val = random.normal(subkey)\n  del subkey  # The subkey is consumed by normal().\n\n  print(f\"draw {i}: {val}\")\n  key = new_key  # new_key is safe to use in the next iteration.\n```\n\nExample:\n```text\ndraw 0: 0.6057640314102173\ndraw 1: -0.21089035272598267\ndraw 2: -0.3948981463909149\n```\n\nExample:\n```text\nkey, subkey = random.split(key)\n```\n\nExample:\n```text\nkey, *forty_two_subkeys = random.split(key, num=43)\n```\n\nExample:\n```text\nkey = random.key(42)\nsubkeys = random.split(key, 3)\nsequence = np.stack([random.normal(subkey) for subkey in subkeys])\nprint(\"individually:\", sequence)\n\nkey = random.key(42)\nprint(\"all at once: \", random.normal(key, shape=(3,)))\n```\n\nExample:\n```text\nindividually: [0.07592554 0.60576403 0.4323065 ]\nall at once:  [-0.02830462  0.46713185  0.29570296]\n```\n\nExample:\n```text\nimport jax\nprint(\"vectorized:\", jax.vmap(random.normal)(subkeys))\n```\n\nExample:\n```text\nvectorized: [0.07592554 0.60576403 0.4323065 ]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.672Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":26,"totalLines":194,"estimatedTokens":913}}10{"id":"doc-key_concepts_jax_documentation-2807fa74","source":"documentation","title":"Key concepts — JAX documentation","url":"https://docs.jax.dev/en/latest/key-concepts.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n\ndef selu(x, alpha=1.67, lambda_=1.05):\n  return lambda_ * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)\n\nselu_jit = jax.jit(selu)\nprint(selu_jit(1.0))\n```\n\nExample:\n```text\n1.05\n```\n\nExample:\n```text\n@jax.jit\ndef selu(x, alpha=1.67, lambda_=1.05):\n  return lambda_ * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  print(x)\n  return x + 1\n\nx = jnp.arange(5)\nresult = f(x)\n```\n\nExample:\n```text\nJitTracer(int32[5])\n```\n\nExample:\n```text\ndef selu(x, alpha=1.67, lambda_=1.05):\n  return lambda_ * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)\n```\n\nExample:\n```text\nx = jnp.arange(5.0)\njax.make_jaxpr(selu)(x)\n```\n\nExample:\n```text\n{ lambda ; a:f32[5]. let\n    b:bool[5] = gt a 0.0:f32[]\n    c:f32[5] = exp a\n    d:f32[5] = mul 1.6699999570846558:f32[] c\n    e:f32[5] = sub d 1.6699999570846558:f32[]\n    f:f32[5] = jit[\n      name=_where\n      jaxpr={ lambda ; b:bool[5] a:f32[5] e:f32[5]. let\n          f:f32[5] = select_n b e a\n        in (f,) }\n    ] b a e\n    g:f32[5] = mul 1.0499999523162842:f32[] f\n  in (g,) }\n```\n\nExample:\n```text\n# (nested) list of parameters\nparams = [1, 2, (jnp.arange(3), jnp.ones(2))]\n\nprint(jax.tree.structure(params))\nprint(jax.tree.leaves(params))\n```\n\nExample:\n```text\nPyTreeDef([*, *, (*, *)])\n[1, 2, Array([0, 1, 2], dtype=int32), Array([1., 1.], dtype=float32)]\n```\n\nExample:\n```text\n# Dictionary of parameters\nparams = {'n': 5, 'W': jnp.ones((2, 2)), 'b': jnp.zeros(2)}\n\nprint(jax.tree.structure(params))\nprint(jax.tree.leaves(params))\n```\n\nExample:\n```text\nPyTreeDef({'W': *, 'b': *, 'n': *})\n[Array([[1., 1.],\n       [1., 1.]], dtype=float32), Array([0., 0.], dtype=float32), 5]\n```\n\nExample:\n```text\n# Named tuple of parameters\nfrom typing import NamedTuple\n\nclass Params(NamedTuple):\n  a: int\n  b: float\n\nparams = Params(1, 5.0)\nprint(jax.tree.structure(params))\nprint(jax.tree.leaves(params))\n```\n\nExample:\n```text\nPyTreeDef(CustomNode(namedtuple[Params], [*, *]))\n[1, 5.0]\n```\n\nExample:\n```text\nimport jax.numpy as jnp\njnp.add(1, 1.0)  # jax.numpy API implicitly promotes mixed types.\n```\n\nExample:\n```text\nArray(2., dtype=float32, weak_type=True)\n```\n\nExample:\n```text\nfrom jax import lax\nlax.add(1, 1.0)  # jax.lax API requires explicit type promotion.\n```\n\nExample:\n```text\n---------------------------------------------------------------------------\nTypeError                                 Traceback (most recent call last)\nCell In[10], line 2\n      1 from jax import lax\n----> 2 lax.add(1, 1.0)  # jax.lax API requires explicit type promotion.\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/lax/lax.py:1165, in add(x, y)\n   1145 r\"\"\"Elementwise addition: :math:`x + y`.\n   1146 \n   1147 This function lowers directly to the `stablehlo.add`_ operation.\n   (...)   1162 .. _stablehlo.add: https://openxla.org/stablehlo/spec#add\n   1163 \"\"\"\n   1164 x, y = core.auto_insert_reshard(x, y)\n-> 1165 return add_p.bind(x, y)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:724, in Primitive.bind(self, *args, **params)\n    722 trace_ctx.set_trace(None)\n    723 try:\n--> 724   return self.bind_with_trace(prev_trace, args, avals, params)\n    725 finally:\n    726   trace_ctx.set_trace(prev_trace)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:732, in Primitive.bind_with_trace(self, trace, args, avals, params)\n    730   with set_current_trace(trace):\n    731     return self.to_lojax(*args, **params)\n--> 732 return trace.process_primitive(self, args, params)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:1304, in EvalTrace.process_primitive(self, primitive, args, params)\n   1302 args = map(full_lower, args)\n   1303 check_eval_args(args)\n-> 1304 return primitive.impl(*args, **params)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/dispatch.py:88, in apply_primitive(prim, *args, **params)\n     86 prev = config.disable_jit.swap_local(False)\n     87 try:\n---> 88   outs = fun(*args)\n     89 finally:\n     90   config.disable_jit.set_local(prev)\n\n    [... skipping hidden 15 frame]\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/lax/lax.py:9650, in check_same_dtypes(name, *avals)\n   9648   equiv = _JNP_FUNCTION_EQUIVALENTS[name]\n   9649   msg += f\" (Tip: jnp.{equiv} is a similar function that does automatic type promotion on inputs).\"\n-> 9650 raise TypeError(msg.format(name, \", \".join(str(a.dtype) for a in avals)))\n\nTypeError: lax.add requires arguments to have the same dtypes, got int32, float32. (Tip: jnp.add is a similar function that does automatic type promotion on inputs).\n```\n\nExample:\n```text\nlax.add(jnp.float32(1), 1.0)\n```\n\nExample:\n```text\nArray(2., dtype=float32)\n```\n\nExample:\n```text\nx = jnp.array([1, 2, 1])\ny = jnp.ones(10)\njnp.convolve(x, y)\n```\n\nExample:\n```text\nArray([1., 3., 4., 4., 4., 4., 4., 4., 4., 4., 3., 1.], dtype=float32)\n```\n\nExample:\n```text\nfrom jax import lax\nresult = lax.conv_general_dilated(\n    x.reshape(1, 1, 3).astype(float),  # note: explicit promotion\n    y.reshape(1, 1, 10),\n    window_strides=(1,),\n    padding=[(len(y) - 1, len(y) - 1)])  # equivalent of padding='full' in NumPy\nresult[0, 0]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.673Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":23,"totalLines":222,"estimatedTokens":1363}}11{"id":"doc-device_local_array_layout_control_jax_documentat-2b09d0c1","source":"documentation","title":"Device-local array layout control — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/layout.html","text":"Example:\n```text\nimport jax, jax.numpy as jnp\nfrom jax.experimental.layout import Layout, Format\nfrom jax.sharding import SingleDeviceSharding\nimport numpy as np\n\ndef init_fn(x, y):\n  return x * 2, y * 3\n\ndef apply_fn(x, y):\n  return x[0, :], y[:, 0]\n```\n\nExample:\n```text\nshape = (4 * 128, 8 * 128)\nduck = jax.ShapeDtypeStruct(shape, jnp.float32)\n\n# Compile the `apply` function with layouts inferred automatically\napply_exe = jax.jit(\n    apply_fn,\n    in_shardings=Format(Layout.AUTO),\n    out_shardings=Format(Layout.AUTO),\n).trace(duck, duck).lower().compile()\n\n# Read back the inferred input layout\narg_formats, kwarg_formats = apply_exe.input_formats\nassert len(kwarg_formats) == 0\nassert arg_formats[0].layout.major_to_minor == (0, 1)\nassert arg_formats[1].layout.major_to_minor == (1, 0)\n```\n\nExample:\n```text\ninit_exe = jax.jit(init_fn, out_shardings=arg_formats).trace(\n    duck, duck).lower().compile()\n\nassert init_exe.output_formats == arg_formats\n```\n\nExample:\n```text\ndef test(x, y, msg):\n  print(f'-- {msg}:')\n  print('x major_to_minor =', x.format.layout.major_to_minor)\n  print('y major_to_minor =', y.format.layout.major_to_minor)\n  try:\n    apply_exe(x, y)\n    print('-> `apply` called successfully')\n  except ValueError as e:\n    assert 'does not match' in str(e)\n    print('-> error: mismatched input layouts')\n  print()\n\ndev = jax.devices()[0]\n\nx1 = y1 = jnp.ones(shape)\ntest(x1, y1, 'uncommitted with mismatched layout')\n\nx2, y2 = init_exe(x1, y1)\ntest(x2, y2, 'uncommitted with matching layout')\n\nx3 = jnp.ones(shape)\ny3 = jax.device_put(np.ones(shape), Format(Layout(major_to_minor=(1, 0)),\n                                           SingleDeviceSharding(dev)))\ntest(x3, y3, 'committed with matching layout')\n\nx4 = jnp.ones(shape)\ny4 = jax.device_put(np.ones(shape), Format(Layout(major_to_minor=(0, 1)),\n                                           SingleDeviceSharding(dev)))\ntest(x4, y4, 'committed with mismatched layout')\n```\n\nExample:\n```text\n-- uncommitted with mismatched layout:\nx major_to_minor = (0, 1)\ny major_to_minor = (0, 1)\n-> `apply` called successfully\n\n-- uncommitted with matching layout:\nx major_to_minor = (0, 1)\ny major_to_minor = (1, 0)\n-> `apply` called successfully\n\n-- committed with matching layout:\nx major_to_minor = (0, 1)\ny major_to_minor = (1, 0)\n-> `apply` called successfully\n\n-- committed with mismatched layout:\nx major_to_minor = (0, 1)\ny major_to_minor = (0, 1)\n-> error: mismatched input layouts\n```\n\nExample:\n```text\nfrom jax.experimental.layout import with_layout_constraint\n\n@jax.jit\ndef f(x):\n  y = x.T\n  # Enforce a specific layout on `y`\n  y = with_layout_constraint(y, Layout(major_to_minor=(0, 1)))\n  return y * 2\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.674Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":6,"totalLines":110,"estimatedTokens":676}}12{"id":"doc-control_flow_and_logical_operators_with_jit_jax_-9480f5e6","source":"documentation","title":"Control flow and logical operators with JIT — JAX documentation","url":"https://docs.jax.dev/en/latest/control-flow.html","text":"Example:\n```text\nfrom jax import grad, jit\nimport jax.numpy as jnp\n```\n\nExample:\n```text\n@jit\ndef f(x):\n  for i in range(3):\n    x = 2 * x\n  return x\n\nprint(f(3))\n```\n\nExample:\n```text\n24\n```\n\nExample:\n```text\n@jit\ndef g(x):\n  y = 0.\n  for i in range(x.shape[0]):\n    y = y + x[i]\n  return y\n\nprint(g(jnp.array([1., 2., 3.])))\n```\n\nExample:\n```text\n6.0\n```\n\nExample:\n```text\n@jit\ndef f(x):\n  if x < 3:\n    return 3. * x ** 2\n  else:\n    return -4 * x\n\n# This will fail!\nf(2)\n```\n\nExample:\n```text\n---------------------------------------------------------------------------\nTracerBoolConversionError                 Traceback (most recent call last)\nCell In[4], line 9\n      5   else:\n      6     return -4 * x\n      7 \n      8 # This will fail!\n----> 9 f(2)\n\n    [... skipping hidden 5 frame]\n\nCell In[4], line 3, in f(x)\n      1 @jit\n      2 def f(x):\n----> 3   if x < 3:\n      4     return 3. * x ** 2\n      5   else:\n      6     return -4 * x\n\n    [... skipping hidden 1 frame]\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:2022, in concretization_function_error.<locals>.error(self, arg)\n   2021 def error(self, arg):\n-> 2022   raise TracerBoolConversionError(arg)\n\nTracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].\nThe error occurred while tracing the function f at /tmp/ipykernel_1621/3402096563.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument x.\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError\n```\n\nExample:\n```text\n@jit\ndef g(x):\n  return (x > 0) and (x < 3)\n\n# This will fail!\ng(2)\n```\n\nExample:\n```text\n---------------------------------------------------------------------------\nTracerBoolConversionError                 Traceback (most recent call last)\nCell In[5], line 6\n      2 def g(x):\n      3   return (x > 0) and (x < 3)\n      4 \n      5 # This will fail!\n----> 6 g(2)\n\n    [... skipping hidden 5 frame]\n\nCell In[5], line 3, in g(x)\n      1 @jit\n      2 def g(x):\n----> 3   return (x > 0) and (x < 3)\n\n    [... skipping hidden 1 frame]\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:2022, in concretization_function_error.<locals>.error(self, arg)\n   2021 def error(self, arg):\n-> 2022   raise TracerBoolConversionError(arg)\n\nTracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].\nThe error occurred while tracing the function g at /tmp/ipykernel_1621/543860509.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument x.\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError\n```\n\nExample:\n```text\ndef f(x):\n  if x < 3:\n    return 3. * x ** 2\n  else:\n    return -4 * x\n\nf = jit(f, static_argnames='x')\n\nprint(f(2.))\n```\n\nExample:\n```text\n12.0\n```\n\nExample:\n```text\ndef f(x, n):\n  y = 0.\n  for i in range(n):\n    y = y + x[i]\n  return y\n\nf = jit(f, static_argnames='n')\n\nf(jnp.array([2., 3., 4.]), 2)\n```\n\nExample:\n```text\nArray(5., dtype=float32)\n```\n\nExample:\n```text\ndef example_fun(length, val):\n  return jnp.ones((length,)) * val\n# un-jit'd works fine\nprint(example_fun(5, 4))\n```\n\nExample:\n```text\n[4. 4. 4. 4. 4.]\n```\n\nExample:\n```text\nbad_example_jit = jit(example_fun)\n# this will fail:\nbad_example_jit(10, 4)\n```\n\nExample:\n```text\n---------------------------------------------------------------------------\nTypeError                                 Traceback (most recent call last)\nCell In[9], line 3\n      1 bad_example_jit = jit(example_fun)\n      2 # this will fail:\n----> 3 bad_example_jit(10, 4)\n\n    [... skipping hidden 5 frame]\n\nCell In[8], line 2, in example_fun(length, val)\n      1 def example_fun(length, val):\n----> 2   return jnp.ones((length,)) * val\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/array_creation.py:139, in ones(shape, dtype, device, out_sharding)\n    137   raise TypeError(\"expected sequence object with len >= 0 or a single integer\")\n    138 if (m := _check_forgot_shape_tuple(\"ones\", shape, dtype)): raise TypeError(m)\n--> 139 shape = canonicalize_shape(shape)\n    140 dtype = dtypes.check_and_canonicalize_user_dtype(\n    141     float if dtype is None else dtype, \"ones\")\n    142 sharding = util.choose_device_or_out_sharding(\n    143     device, out_sharding, 'jnp.ones')\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/array_creation.py:46, in canonicalize_shape(shape, context)\n     44   return core.canonicalize_shape((shape,), context)\n     45 else:\n---> 46   return core.canonicalize_shape(shape, context)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:2125, in canonicalize_shape(shape, context)\n   2123 except TypeError:\n   2124   pass\n-> 2125 raise _invalid_shape_error(shape, context)\n\nTypeError: Shapes must be 1D sequences of concrete values of integer type, got (JitTracer(~int32[]),).\nIf using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.\nThe error occurred while tracing the function example_fun at /tmp/ipykernel_1621/1210496444.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument length.\n```\n\nExample:\n```text\n# static_argnames tells JAX to recompile on changes at these argument positions:\ngood_example_jit = jit(example_fun, static_argnames='length')\n# first compile\nprint(good_example_jit(10, 4))\n# recompiles\nprint(good_example_jit(5, 4))\n```\n\nExample:\n```text\n[4. 4. 4. 4. 4. 4. 4. 4. 4. 4.]\n[4. 4. 4. 4. 4.]\n```\n\nExample:\n```text\n@jit\ndef f(x):\n  print(x)\n  y = 2 * x\n  print(y)\n  return y\nf(2)\n```\n\nExample:\n```text\nJitTracer(~int32[])\nJitTracer(~int32[])\n```\n\nExample:\n```text\nArray(4, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\ndef cond(pred, true_fun, false_fun, operand):\n  if pred:\n    return true_fun(operand)\n  else:\n    return false_fun(operand)\n```\n\nExample:\n```text\nfrom jax import lax\n\noperand = jnp.array([0.])\nlax.cond(True, lambda x: x+1, lambda x: x-1, operand)\n# --> array([1.], dtype=float32)\nlax.cond(False, lambda x: x+1, lambda x: x-1, operand)\n# --> array([-1.], dtype=float32)\n```\n\nExample:\n```text\nArray([-1.], dtype=float32)\n```\n\nExample:\n```text\ndef while_loop(cond_fun, body_fun, init_val):\n  val = init_val\n  while cond_fun(val):\n    val = body_fun(val)\n  return val\n```\n\nExample:\n```text\ninit_val = 0\ncond_fun = lambda x: x < 10\nbody_fun = lambda x: x+1\nlax.while_loop(cond_fun, body_fun, init_val)\n# --> array(10, dtype=int32)\n```\n\nExample:\n```text\nArray(10, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\ndef fori_loop(start, stop, body_fun, init_val):\n  val = init_val\n  for i in range(start, stop):\n    val = body_fun(i, val)\n  return val\n```\n\nExample:\n```text\ninit_val = 0\nstart = 0\nstop = 10\nbody_fun = lambda i,x: x+i\nlax.fori_loop(start, stop, body_fun, init_val)\n# --> array(45, dtype=int32)\n```\n\nExample:\n```text\nArray(45, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\ndef python_check_positive_even(x):\n  is_even = x % 2 == 0\n  # `and` short-circults, so when `is_even` is `False`, `x > 0` is not evaluated.\n  return is_even and (x > 0)\n\n@jit\ndef jax_check_positive_even(x):\n  is_even = x % 2 == 0\n  # `logical_and` does not short circuit, so `x > 0` is always evaluated.\n  return jnp.logical_and(is_even, x > 0)\n\nprint(python_check_positive_even(24))\nprint(jax_check_positive_even(24))\n```\n\nExample:\n```text\nTrue\nTrue\n```\n\nExample:\n```text\nx = jnp.array([-1, 2, 5])\nprint(jax_check_positive_even(x))\n```\n\nExample:\n```text\n[False  True False]\n```\n\nExample:\n```text\nprint(python_check_positive_even(x))\n```\n\nExample:\n```text\n---------------------------------------------------------------------------\nValueError                                Traceback (most recent call last)\nCell In[17], line 1\n----> 1 print(python_check_positive_even(x))\n\nCell In[15], line 4, in python_check_positive_even(x)\n      1 def python_check_positive_even(x):\n      2   is_even = x % 2 == 0\n      3   # `and` short-circults, so when `is_even` is `False`, `x > 0` is not evaluated.\n----> 4   return is_even and (x > 0)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/array.py:295, in ArrayImpl.__bool__(self)\n    294 def __bool__(self):\n--> 295   core.check_bool_conversion(self)\n    296   return bool(self._value)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:941, in check_bool_conversion(arr)\n    938   raise ValueError(\"The truth value of an empty array is ambiguous. Use\"\n    939                    \" `array.size > 0` to check that an array is not empty.\")\n    940 if arr.size > 1:\n--> 941   raise ValueError(\"The truth value of an array with more than one element\"\n    942                    \" is ambiguous. Use a.any() or a.all()\")\n\nValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()\n```\n\nExample:\n```text\ndef f(x):\n  if x < 3:\n    return 3. * x ** 2\n  else:\n    return -4 * x\n\nprint(grad(f)(2.))  # ok!\nprint(grad(f)(4.))  # ok!\n```\n\nExample:\n```text\n12.0\n-4.0\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.675Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":39,"totalLines":414,"estimatedTokens":2332}}13{"id":"doc-complex_numbers_and_differentiation_jax_document-b880c229","source":"documentation","title":"Complex numbers and differentiation — JAX documentation","url":"https://docs.jax.dev/en/latest/complex-differentiation.html","text":"Example:\n```text\nimport jax.numpy as jnp\n\ndef f(z):\n  x, y = jnp.real(z), jnp.imag(z)\n  return u(x, y) + v(x, y) * 1j\n\ndef g(x, y):\n  return (u(x, y), v(x, y))\n```\n\nExample:\n```text\nfrom jax import random, grad, jvp\n\ndef check(seed):\n  key = random.key(seed)\n\n  # random coeffs for u and v\n  key, subkey = random.split(key)\n  a, b, c, d = random.uniform(subkey, (4,))\n\n  def fun(z):\n    x, y = jnp.real(z), jnp.imag(z)\n    return u(x, y) + v(x, y) * 1j\n\n  def u(x, y):\n    return a * x + b * y\n\n  def v(x, y):\n    return c * x + d * y\n\n  # primal point\n  key, subkey = random.split(key)\n  x, y = random.uniform(subkey, (2,))\n  z = x + y * 1j\n\n  # tangent vector\n  key, subkey = random.split(key)\n  c, d = random.uniform(subkey, (2,))\n  z_dot = c + d * 1j\n\n  # check jvp\n  _, ans = jvp(fun, (z,), (z_dot,))\n  expected = (grad(u, 0)(x, y) * c +\n              grad(u, 1)(x, y) * d +\n              grad(v, 0)(x, y) * c * 1j+\n              grad(v, 1)(x, y) * d * 1j)\n  print(jnp.allclose(ans, expected))\n```\n\nExample:\n```text\ncheck(0)\ncheck(1)\ncheck(2)\n```\n\nExample:\n```text\nTrue\nTrue\nTrue\n```\n\nExample:\n```text\nfrom jax import vjp\n\ndef check(seed):\n  key = random.key(seed)\n\n  # random coeffs for u and v\n  key, subkey = random.split(key)\n  a, b, c, d = random.uniform(subkey, (4,))\n\n  def fun(z):\n    x, y = jnp.real(z), jnp.imag(z)\n    return u(x, y) + v(x, y) * 1j\n\n  def u(x, y):\n    return a * x + b * y\n\n  def v(x, y):\n    return c * x + d * y\n\n  # primal point\n  key, subkey = random.split(key)\n  x, y = random.uniform(subkey, (2,))\n  z = x + y * 1j\n\n  # cotangent vector\n  key, subkey = random.split(key)\n  c, d = random.uniform(subkey, (2,))\n  z_bar = jnp.array(c + d * 1j)  # for dtype control\n\n  # check vjp\n  _, fun_vjp = vjp(fun, z)\n  ans, = fun_vjp(z_bar)\n  expected = (grad(u, 0)(x, y) * c +\n              grad(v, 0)(x, y) * (-d) +\n              grad(u, 1)(x, y) * c * (-1j) +\n              grad(v, 1)(x, y) * (-d) * (-1j))\n  assert jnp.allclose(ans, expected, atol=1e-5, rtol=1e-5)\n```\n\nExample:\n```text\ndef f(z):\n  x, y = jnp.real(z), jnp.imag(z)\n  return x**2 + y**2\n\nz = 3. + 4j\ngrad(f)(z)\n```\n\nExample:\n```text\nArray(6.-8.j, dtype=complex64)\n```\n\nExample:\n```text\ndef f(z):\n  return jnp.sin(z)\n\nz = 3. + 4j\ngrad(f, holomorphic=True)(z)\n```\n\nExample:\n```text\nArray(-27.034946-3.8511534j, dtype=complex64, weak_type=True)\n```\n\nExample:\n```text\ndef f(z):\n  return jnp.conjugate(z)\n\nz = 3. + 4j\ngrad(f, holomorphic=True)(z)  # f is not actually holomorphic!\n```\n\nExample:\n```text\nArray(1.-0.j, dtype=complex64, weak_type=True)\n```\n\nExample:\n```text\nA = jnp.array([[5.,    2.+3j,    5j],\n              [2.-3j,   7.,  1.+7j],\n              [-5j,  1.-7j,    12.]])\n\ndef f(X):\n    L = jnp.linalg.cholesky(X)\n    return jnp.sum((L - jnp.sin(L))**2)\n\ngrad(f, holomorphic=True)(A)\n```\n\nExample:\n```text\nArray([[-0.7534186  +0.j       , -3.0509028 -10.940544j ,\n         5.9896846  +3.5423026j],\n       [-3.0509028 +10.940544j , -8.904491   +0.j       ,\n        -5.1351523  -6.559373j ],\n       [ 5.9896846  -3.5423026j, -5.1351523  +6.559373j ,\n         0.01320427 +0.j       ]], dtype=complex64)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.676Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":13,"totalLines":174,"estimatedTokens":780}}14{"id":"doc-distributed_arrays_and_automatic_parallelization-7427b1d9","source":"documentation","title":"Distributed arrays and automatic parallelization — JAX documentation","url":"https://docs.jax.dev/en/latest/parallel.html","text":"Example:\n```text\nfrom __future__ import annotations\nimport enum\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\njax.config.update('jax_num_cpu_devices', 8)\n```\n\nExample:\n```text\njax.set_mesh(jax.make_mesh((4, 2), ('X', 'Y')))  # explicit mode by default\n\nx = jnp.arange(8 * 4.).reshape(8, 4)\nx = jax.device_put(x, jax.P('X', 'Y'))\nprint(jax.typeof(x))  # f32[8@X, 4@Y]\n```\n\nExample:\n```text\nfloat32[8@X,4@Y]\n```\n\nExample:\n```text\njax.debug.visualize_array_sharding(x)\n```\n\nExample:\n```text\nCPU 0    CPU 1  \n                  \n                  \n  CPU 2    CPU 3  \n                  \n                  \n  CPU 4    CPU 5  \n                  \n                  \n  CPU 6    CPU 7\n```\n\nExample:\n```text\ny = jnp.sin(x).T\nprint(jax.typeof(y))  # f32[4@Y, 8@X]\n```\n\nExample:\n```text\nfloat32[4@Y,8@X]\n```\n\nExample:\n```text\nclass AbstractMesh:\n  axis_sizes: tuple[int, ...]\n  axis_names: tuple[str, ...]\n  axis_types: tuple[AxisType, ...]\n\nclass AxisType(enum.Enum):\n  Auto = enum.auto()\n  Explicit = enum.auto()\n  Manual = enum.auto()\n\n# A concrete mesh additionally includes physical device objects with e.g.\n# precise coordinates:\n```\n\nExample:\n```text\nimport numpy as np\n\nclass Mesh:\n  devices: np.ndarray[jax.Device]\n  axis_names: tuple[str, ...]\n  axis_types: tuple[AxisType, ...]\n\n  @property\n  def axis_sizes(self) -> tuple[int, ...]:\n    return self.devices.shape\n```\n\nExample:\n```text\nmesh = jax.make_mesh((4, 2), ('X', 'Y'))\nprint(mesh)\n```\n\nExample:\n```text\nMesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit))\n```\n\nExample:\n```text\njax.set_mesh(mesh)\n```\n\nExample:\n```text\n<jax._src.sharding_impls.set_mesh at 0x71185e5b3610>\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  abstract_mesh = jax.sharding.AbstractMesh((8,), ('A',), (jax.sharding.AxisType.Explicit,))\n  with jax.sharding.use_abstract_mesh(abstract_mesh):\n    y = jax.reshard(x, jax.P('A', None))\n    return y * 2\n\nz = f(x)\nprint(jax.typeof(z))  # f32[8@A, 4]\n```\n\nExample:\n```text\nfloat32[8@A,4]\n```\n\nExample:\n```text\nprint(x.sharding)\njax.debug.visualize_array_sharding(x)\n```\n\nExample:\n```text\nNamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P('X', 'Y'), memory_kind=device)\n```\n\nExample:\n```text\nfor s in x.addressable_shards:\n  print(s.device, s.data, sep='\\n', end='\\n\\n')\n```\n\nExample:\n```text\ncpu:0\n[[0. 1.]\n [4. 5.]]\n\ncpu:1\n[[2. 3.]\n [6. 7.]]\n\ncpu:2\n[[ 8.  9.]\n [12. 13.]]\n\ncpu:3\n[[10. 11.]\n [14. 15.]]\n\ncpu:4\n[[16. 17.]\n [20. 21.]]\n\ncpu:5\n[[18. 19.]\n [22. 23.]]\n\ncpu:6\n[[24. 25.]\n [28. 29.]]\n\ncpu:7\n[[26. 27.]\n [30. 31.]]\n```\n\nExample:\n```text\ny = jax.device_put(x, jax.P('Y', 'X'))\nprint(y.sharding)\njax.debug.visualize_array_sharding(y)\n```\n\nExample:\n```text\nNamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P('Y', 'X'), memory_kind=device)\n```\n\nExample:\n```text\nCPU 0    CPU 2    CPU 4    CPU 6  \n                                    \n                                    \n                                    \n                                    \n                                    \n  CPU 1    CPU 3    CPU 5    CPU 7\n```\n\nExample:\n```text\ny = jax.device_put(x, jax.P('X', None))\nprint(y.sharding)\njax.debug.visualize_array_sharding(y)\n```\n\nExample:\n```text\nNamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P('X', None), memory_kind=device)\n```\n\nExample:\n```text\nCPU 0,1   \n            \n            \n  CPU 2,3   \n            \n            \n  CPU 4,5   \n            \n            \n  CPU 6,7\n```\n\nExample:\n```text\nfor s in y.addressable_shards:\n  print(s.device, s.data, sep='\\n', end='\\n\\n')\n```\n\nExample:\n```text\ncpu:0\n[[0. 1. 2. 3.]\n [4. 5. 6. 7.]]\n\ncpu:1\n[[0. 1. 2. 3.]\n [4. 5. 6. 7.]]\n\ncpu:2\n[[ 8.  9. 10. 11.]\n [12. 13. 14. 15.]]\n\ncpu:3\n[[ 8.  9. 10. 11.]\n [12. 13. 14. 15.]]\n\ncpu:4\n[[16. 17. 18. 19.]\n [20. 21. 22. 23.]]\n\ncpu:5\n[[16. 17. 18. 19.]\n [20. 21. 22. 23.]]\n\ncpu:6\n[[24. 25. 26. 27.]\n [28. 29. 30. 31.]]\n\ncpu:7\n[[24. 25. 26. 27.]\n [28. 29. 30. 31.]]\n```\n\nExample:\n```text\ny = jax.device_put(x, jax.P(('X', 'Y')))\nprint(y.sharding)\njax.debug.visualize_array_sharding(y)\n```\n\nExample:\n```text\nNamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P(('X', 'Y'),), memory_kind=device)\n```\n\nExample:\n```text\ny = jax.device_put(x, jax.P('X', None, unreduced={'Y'}))\nprint(y.sharding)\n```\n\nExample:\n```text\nNamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P('X', None, unreduced={'Y'}, unreduced_kind=sum), memory_kind=device)\n```\n\nExample:\n```text\ncpu:0\n[[0. 1. 0. 0.]\n [4. 5. 0. 0.]]\n\ncpu:1\n[[0. 0. 2. 3.]\n [0. 0. 6. 7.]]\n\ncpu:2\n[[ 8.  9.  0.  0.]\n [12. 13.  0.  0.]]\n\ncpu:3\n[[ 0.  0. 10. 11.]\n [ 0.  0. 14. 15.]]\n\ncpu:4\n[[16. 17.  0.  0.]\n [20. 21.  0.  0.]]\n\ncpu:5\n[[ 0.  0. 18. 19.]\n [ 0.  0. 22. 23.]]\n\ncpu:6\n[[24. 25.  0.  0.]\n [28. 29.  0.  0.]]\n\ncpu:7\n[[ 0.  0. 26. 27.]\n [ 0.  0. 30. 31.]]\n```\n\nExample:\n```text\nmesh2 = jax.make_mesh((8,), ('A',))\nz = jax.device_put(x, jax.NamedSharding(mesh2, jax.P('A', None)))\nprint(z.sharding)\nprint(y.sharding)\n```\n\nExample:\n```text\nNamedSharding(mesh=Mesh('A': 8, axis_types=(Explicit,)), spec=P('A', None), memory_kind=device)\nNamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit)), spec=P('X', None, unreduced={'Y'}, unreduced_kind=sum), memory_kind=device)\n```\n\nExample:\n```text\nprint(jax.typeof(x).sharding)\n```\n\nExample:\n```text\nNamedSharding(mesh=AbstractMesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit), device_kind=cpu, num_cores=None, platform=cpu), spec=P('X', 'Y'))\n```\n\nExample:\n```text\njax.jit(lambda x: print(jax.typeof(x).sharding))(x)\n```\n\nExample:\n```text\n<array_type> ::= <dtype>[<size_and_sharding>, ...]\n <size_and_sharding> ::= <size> | <size>@<MeshAxisName>\n```\n\nExample:\n```text\narg0 = jax.device_put(np.arange(4).reshape(4, 1), jax.P(\"X\", None))\narg1 = jax.device_put(np.arange(8).reshape(1, 8), jax.P(None, \"Y\"))\n\nresult = arg0 + arg1\n\nprint(f\"{jax.typeof(arg0)=!s}\")\nprint(f\"{jax.typeof(arg1)=!s}\")\nprint(f\"{jax.typeof(result)=!s}\")\n```\n\nExample:\n```text\njax.typeof(arg0)=int32[4@X,1]\njax.typeof(arg1)=int32[1,8@Y]\njax.typeof(result)=int32[4@X,8@Y]\n```\n\nExample:\n```text\n@jax.jit\ndef add_arrays(x, y):\n  ans = x + y\n  print(f\"{jax.typeof(arg0)=!s}\")\n  print(f\"{jax.typeof(arg1)=!s}\")\n  print(f\"{jax.typeof(result)=!s}\")\n  return ans\n\nadd_arrays(arg0, arg1)\n```\n\nExample:\n```text\nArray([[ 0,  1,  2,  3,  4,  5,  6,  7],\n       [ 1,  2,  3,  4,  5,  6,  7,  8],\n       [ 2,  3,  4,  5,  6,  7,  8,  9],\n       [ 3,  4,  5,  6,  7,  8,  9, 10]], dtype=int32)\n```\n\nExample:\n```text\nx = jax.random.normal(jax.random.key(0), (8, 4),\n                      out_sharding=jax.P('X', 'Y'))\nprint(jax.typeof(x))\n```\n\nExample:\n```text\ny = x.sum(0)\nprint(jax.typeof(y))\n```\n\nExample:\n```text\nfloat32[4@Y]\n```\n\nExample:\n```text\ncompile_txt = jax.jit(lambda x: x.sum(0)).lower(x).compile().as_text()\nprint('all-reduce(' in compile_txt)\n```\n\nExample:\n```text\nTrue\n```\n\nExample:\n```text\nx = jax.device_put(jnp.arange(8 * 4.).reshape(8, 4), jax.P(None, 'X'))\ny = jax.device_put(jnp.arange(4 * 16.).reshape(4, 16), jax.P('X', None))\n\ntry:\n  jnp.dot(x, y)\nexcept Exception as e:\n  print(\"ERROR!\")\n  print(e)\n```\n\nExample:\n```text\nERROR!\nContracting dimensions are sharded and it is ambiguous how the output should be sharded. Please specify the output sharding via the `out_sharding` parameter. Got lhs_contracting_spec=('X',) and rhs_contracting_spec=('X',)\n```\n\nExample:\n```text\nz = jnp.dot(x, y, out_sharding=jax.P('X', None))\n\nprint(jax.typeof(z))\n```\n\nExample:\n```text\nfloat32[8@X,16]\n```\n\nExample:\n```text\nfrom jax.sharding import auto_axes, explicit_axes\n\nx = jax.device_put(np.arange(16).reshape(4, 4), jax.P(\"X\", None))\ny = jax.device_put(np.arange(16).reshape(4, 4), jax.P(None, \"X\"))\n\ntry:\n  x + y\nexcept Exception as e:\n  print(\"ERROR!\")\n  print(e)\n```\n\nExample:\n```text\nERROR!\nadd operation with inputs: i32[4@X,4], i32[4,4@X] produces an illegally sharded result: i32[4@X,4@X]\n```\n\nExample:\n```text\n@auto_axes\ndef add2(x, y):\n  print(\"We're in auto-sharding mode here. This is the current mesh:\\n\"\n        f\"{jax.sharding.get_abstract_mesh()}\")\n  return x + y\n\nresult = add2(x, y, out_sharding=jax.P(\"X\", None))\nprint(f\"Result type: {jax.typeof(result)}\")\n```\n\nExample:\n```text\nWe're in auto-sharding mode here. This is the current mesh:\nAbstractMesh('X': 4, 'Y': 2, axis_types=(Auto, Auto), device_kind=cpu, num_cores=None, platform=cpu)\nResult type: int32[4@X,4]\n```\n\nExample:\n```text\nAuto = jax.sharding.AxisType.Auto\nauto_mesh = jax.make_mesh((4, 2), ('X', 'Y'), (Auto, Auto))\njax.set_mesh(auto_mesh)\n\nx = jax.device_put(jnp.arange(8 * 4. ).reshape(8, 4 ), jax.P(None, 'X'))\ny = jax.device_put(jnp.arange(4 * 16.).reshape(4, 16), jax.P('X', None))\n\nz = jnp.dot(x, y)  # not an error!\n```\n\nExample:\n```text\nprint(z.sharding)  # works at the top-level only (i.e. outside `jit`)\n```\n\nExample:\n```text\nNamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Auto, Auto)), spec=P(), memory_kind=device)\n```\n\nExample:\n```text\n@jax.jit\ndef f(x, y):\n  z = jnp.dot(x, y)\n  z = jax.lax.with_sharding_constraint(z, jax.P('X', None))\n  return z\n\nz = f(x, y)\nprint(z.sharding)\n```\n\nExample:\n```text\nNamedSharding(mesh=Mesh('X': 4, 'Y': 2, axis_types=(Auto, Auto)), spec=P('X',), memory_kind=device)\n```\n\nExample:\n```text\n@explicit_axes\ndef explicit_g(y):\n  print(f'mesh inside g: {jax.sharding.get_abstract_mesh()}')\n  print(f'y.sharding inside g: {jax.typeof(y) = }')\n  z = y * 2\n  print(f'z.sharding inside g: {jax.typeof(z) = }', end='\\n\\n')\n  return z\n\n@jax.jit\ndef f(arr1):\n  print(f'mesh inside f: {jax.sharding.get_abstract_mesh()}', end='\\n\\n')\n  x = jnp.sin(arr1)\n  z = explicit_g(x, in_sharding=jax.P(\"X\", \"Y\"))\n  return z + 1\n\nx = jax.device_put(np.arange(16).reshape(4, 4), jax.P(\"X\", \"Y\"))\nf(x)\n```\n\nExample:\n```text\nmesh inside f: AbstractMesh('X': 4, 'Y': 2, axis_types=(Auto, Auto), device_kind=cpu, num_cores=None, platform=cpu)\n\n\nmesh inside g: AbstractMesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit), device_kind=cpu, num_cores=None, platform=cpu)\ny.sharding inside g: jax.typeof(y) = ShapedArray(float32[4@X,4@Y])\nz.sharding inside g: jax.typeof(z) = ShapedArray(float32[4@X,4@Y])\n```\n\nExample:\n```text\nArray([[ 1.        ,  2.682942  ,  2.818595  ,  1.28224   ],\n       [-0.513605  , -0.9178486 ,  0.44116902,  2.3139732 ],\n       [ 2.9787164 ,  1.824237  , -0.08804226, -0.99998045],\n       [-0.07314587,  1.840334  ,  2.9812148 ,  2.3005757 ]],      dtype=float32)\n```\n\nExample:\n```text\njax.set_mesh(jax.make_mesh((4, 2), ('X', 'Y')))  # Explicit mode\n\ndef compare_shardings(x):\n  print(f\"=== with mesh: {jax.sharding.get_abstract_mesh()} ===\")\n  print(f\"Concrete value sharding: {x.sharding.spec}\")\n  print(f\"Type-specified sharding: {jax.typeof(x).sharding.spec}\\n\")\n\nmy_array = jnp.sin(jax.device_put(np.arange(8), jax.P(\"X\")))\ncompare_shardings(my_array)\n\n@auto_axes\ndef check_in_auto_context(x):\n  compare_shardings(x)\n  return x\n\ncheck_in_auto_context(my_array, out_sharding=jax.P(\"X\"))\n```\n\nExample:\n```text\n=== with mesh: AbstractMesh('X': 4, 'Y': 2, axis_types=(Explicit, Explicit), device_kind=cpu, num_cores=None, platform=cpu) ===\nConcrete value sharding: P('X',)\nType-specified sharding: P('X',)\n\n=== with mesh: AbstractMesh('X': 4, 'Y': 2, axis_types=(Auto, Auto), device_kind=cpu, num_cores=None, platform=cpu) ===\nConcrete value sharding: P('X',)\nType-specified sharding: P(None,)\n```\n\nExample:\n```text\nArray([ 0.        ,  0.84147096,  0.9092974 ,  0.14112   , -0.7568025 ,\n       -0.9589243 , -0.2794155 ,  0.6569866 ], dtype=float32)\n```\n\nExample:\n```text\nmesh = jax.make_mesh((4, 2), ('X', 'Y'))\njax.set_mesh(mesh)\n\nx = jax.device_put(jnp.arange(8 * 4. ).reshape(8, 4 ), jax.P(None, 'X'))\ny = jax.device_put(jnp.arange(4 * 16.).reshape(4, 16), jax.P('X', None))\n\n@jax.shard_map(out_specs=jax.P('X', None))\ndef matmul(x_shard, y_shard):\n  z_summand = jnp.dot(x_shard, y_shard)\n  return jax.lax.psum_scatter(z_summand, 'X', tiled=True)\n\nz = matmul(x, y)\nprint(jax.typeof(z))\n\nz_ref = jnp.dot(x, y, out_sharding=jax.P('X', None))\nprint(jnp.allclose(z_ref, z))\n```\n\nExample:\n```text\nfloat32[8@X,16]\nTrue\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.678Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":69,"totalLines":641,"estimatedTokens":3020}}15{"id":"doc-colocated_python_jax_documentation-49db644d","source":"documentation","title":"Colocated Python — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/colocated-python.html","text":"Example:\n```text\nimport jax\nimport jax.experimental.colocated_python as colocated_python\n\ndevices = jax.devices()\ncpu_devices = colocated_python.colocated_cpu_devices(devices)\nprint(cpu_devices)\n```\n\nExample:\n```text\n[CpuDevice(id=0)]\n```\n\nExample:\n```text\ncpu_mesh = jax.sharding.Mesh(cpu_devices, [\"x\"])\ncpu_sharding = jax.sharding.NamedSharding(cpu_mesh, jax.P())\nx = jax.device_put(1, cpu_sharding)\ny = jax.jit(lambda x: x + 1)(x)\nprint(y)\n```\n\nExample:\n```text\n2\n```\n\nExample:\n```text\ndef f(x):\n  return x + 1\n\n\nf = colocated_python.colocated_python(f)\ny = f(x)\nassert y.sharding == x.sharding\nprint(y)\n```\n\nExample:\n```text\ndef f(x):\n  with open('/tmp/foo', 'w') as f:\n    f.write(str(x))\n  return x\n\n\nf = colocated_python.colocated_python(f)\njax.block_until_ready(f(x))\n```\n\nExample:\n```text\nArray(1, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\ndef f(x):\n  return x + 1\n\n\nf = colocated_python.colocated_python(f)\nf = f.specialize(out_specs_fn=lambda x: x)\ny = f(x)\nassert y.sharding == x.sharding\n```\n\nExample:\n```text\nimport jax.numpy as jnp\n\n\ndef f(x):\n  return x + 1\n\n\nf = colocated_python.colocated_python(f)\nf = f.specialize(\n    in_specs=(\n        # args\n        (\n            jax.ShapeDtypeStruct(\n                shape=(), dtype=jnp.int32, sharding=cpu_sharding\n            ),\n        ),\n        # kwargs\n        {},\n    ),\n    out_specs_fn=lambda x: jax.ShapeDtypeStruct(\n        shape=(), dtype=jnp.int32, sharding=cpu_sharding\n    ),\n)\nf(x)  # `x` must match the input spec.\n```\n\nExample:\n```text\nArray(2, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\ndef f():\n  with open('/tmp/foo', 'w') as f:\n    f.write('foo')\n  return\n\n\nf = colocated_python.colocated_python(f)\nf = f.specialize(devices=cpu_devices)\nf()  # Would be an error if `f` is not specialized with ``devices``.\n```\n\nExample:\n```text\nclass Adder:\n\n  def __init__(self, increment):\n    print('Adder created')\n    self.increment = increment\n\n  def __del__(self):\n    print('Adder destroyed')\n\n  def add(self, x):\n    return x + self.increment\n\n\nAdder = colocated_python.colocated_python_class(Adder)\nadder = Adder(1)\nx = jax.device_put(1, cpu_sharding)\ny = adder.add(x)\nprint(y)\n```\n\nExample:\n```text\nAdder created\n2\n```\n\nExample:\n```text\ndel adder\n```\n\nExample:\n```text\nAdder destroyed\n```\n\nExample:\n```text\nimport concurrent.futures\nimport time\n\n\ndef f(x):\n  time.sleep(1)\n  return x + 1\n\n\nf = colocated_python.colocated_python(f)\nf = f.specialize(out_specs_fn=lambda x: x)  # Calls will be asynchronous.\n\nwith concurrent.futures.ThreadPoolExecutor(2) as executor:\n  fut1 = executor.submit(f, x)\n  fut2 = executor.submit(f, x)\n  # Will finish in approximately 1 second instead of 2 seconds.\n  jax.block_until_ready([fut1.result(), fut2.result()])\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.678Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":16,"totalLines":176,"estimatedTokens":689}}16{"id":"doc-automatic_differentiation_jax_documentation-471e1cec","source":"documentation","title":"Automatic differentiation — JAX documentation","url":"https://docs.jax.dev/en/latest/automatic-differentiation.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax import grad\n\ngrad_tanh = grad(jnp.tanh)\nprint(grad_tanh(2.0))\n```\n\nExample:\n```text\n0.070650816\n```\n\nExample:\n```text\nprint(grad(grad(jnp.tanh))(2.0))\nprint(grad(grad(grad(jnp.tanh)))(2.0))\n```\n\nExample:\n```text\n-0.13621868\n0.25265405\n```\n\nExample:\n```text\nf = lambda x: x**3 + 2*x**2 - 3*x + 1\n\ndfdx = jax.grad(f)\n```\n\nExample:\n```text\nd2fdx = jax.grad(dfdx)\nd3fdx = jax.grad(d2fdx)\nd4fdx = jax.grad(d3fdx)\n```\n\nExample:\n```text\nprint(dfdx(1.))\nprint(d2fdx(1.))\nprint(d3fdx(1.))\nprint(d4fdx(1.))\n```\n\nExample:\n```text\n4.0\n10.0\n6.0\n0.0\n```\n\nExample:\n```text\nkey = jax.random.key(0)\n\ndef sigmoid(x):\n  return 0.5 * (jnp.tanh(x / 2) + 1)\n\n# Outputs probability of a label being true.\ndef predict(W, b, inputs):\n  return sigmoid(jnp.dot(inputs, W) + b)\n\n# Build a toy dataset.\ninputs = jnp.array([[0.52, 1.12,  0.77],\n                    [0.88, -1.08, 0.15],\n                    [0.52, 0.06, -1.30],\n                    [0.74, -2.49, 1.39]])\ntargets = jnp.array([True, True, False, True])\n\n# Training loss is the negative log-likelihood of the training examples.\ndef loss(W, b):\n  preds = predict(W, b, inputs)\n  label_probs = preds * targets + (1 - preds) * (1 - targets)\n  return -jnp.sum(jnp.log(label_probs))\n\n# Initialize random model coefficients\nkey, W_key, b_key = jax.random.split(key, 3)\nW = jax.random.normal(W_key, (3,))\nb = jax.random.normal(b_key, ())\n```\n\nExample:\n```text\n# Differentiate `loss` with respect to the first positional argument:\nW_grad = grad(loss, argnums=0)(W, b)\nprint(f'{W_grad=}')\n\n# Since argnums=0 is the default, this does the same thing:\nW_grad = grad(loss)(W, b)\nprint(f'{W_grad=}')\n\n# But you can choose different values too, and drop the keyword:\nb_grad = grad(loss, 1)(W, b)\nprint(f'{b_grad=}')\n\n# Including tuple values\nW_grad, b_grad = grad(loss, (0, 1))(W, b)\nprint(f'{W_grad=}')\nprint(f'{b_grad=}')\n```\n\nExample:\n```text\nW_grad=Array([-0.433146 , -0.7354605, -1.2598922], dtype=float32)\nW_grad=Array([-0.433146 , -0.7354605, -1.2598922], dtype=float32)\nb_grad=Array(-0.69001776, dtype=float32)\nW_grad=Array([-0.433146 , -0.7354605, -1.2598922], dtype=float32)\nb_grad=Array(-0.69001776, dtype=float32)\n```\n\nExample:\n```text\ndef loss2(params_dict):\n    preds = predict(params_dict['W'], params_dict['b'], inputs)\n    label_probs = preds * targets + (1 - preds) * (1 - targets)\n    return -jnp.sum(jnp.log(label_probs))\n\nprint(grad(loss2)({'W': W, 'b': b}))\n```\n\nExample:\n```text\n{'W': Array([-0.433146 , -0.7354605, -1.2598922], dtype=float32), 'b': Array(-0.69001776, dtype=float32)}\n```\n\nExample:\n```text\nloss_value, Wb_grad = jax.value_and_grad(loss, (0, 1))(W, b)\nprint('loss value', loss_value)\nprint('loss value', loss(W, b))\n```\n\nExample:\n```text\nloss value 2.9729187\nloss value 2.9729187\n```\n\nExample:\n```text\n# Set a step size for finite differences calculations\neps = 1e-4\n\n# Check b_grad with scalar finite differences\nb_grad_numerical = (loss(W, b + eps / 2.) - loss(W, b - eps / 2.)) / eps\nprint('b_grad_numerical', b_grad_numerical)\nprint('b_grad_autodiff', grad(loss, 1)(W, b))\n\n# Check W_grad with finite differences in a random direction\nkey, subkey = jax.random.split(key)\nvec = jax.random.normal(subkey, W.shape)\nunitvec = vec / jnp.sqrt(jnp.vdot(vec, vec))\nW_grad_numerical = (loss(W + eps / 2. * unitvec, b) - loss(W - eps / 2. * unitvec, b)) / eps\nprint('W_dirderiv_numerical', W_grad_numerical)\nprint('W_dirderiv_autodiff', jnp.vdot(grad(loss)(W, b), unitvec))\n```\n\nExample:\n```text\nb_grad_numerical -0.6890297\nb_grad_autodiff -0.69001776\nW_dirderiv_numerical 1.3041496\nW_dirderiv_autodiff 1.3006744\n```\n\nExample:\n```text\nfrom jax.test_util import check_grads\n\ncheck_grads(loss, (W, b), order=2)  # check up to 2nd order derivatives\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.680Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":18,"totalLines":179,"estimatedTokens":944}}17{"id":"doc-introduction_to_debugging_jax_documentation-9f2a0152","source":"documentation","title":"Introduction to debugging — JAX documentation","url":"https://docs.jax.dev/en/latest/debugging.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n\n@jax.jit\ndef f(x):\n  print(\"print(x) ->\", x)\n  y = jnp.sin(x)\n  print(\"print(y) ->\", y)\n  return y\n\nresult = f(2.)\n```\n\nExample:\n```text\nprint(x) -> JitTracer(~float32[])\nprint(y) -> JitTracer(~float32[])\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  jax.debug.print(\"jax.debug.print(x) -> {x}\", x=x)\n  y = jnp.sin(x)\n  jax.debug.print(\"jax.debug.print(y) -> {y}\", y=y)\n  return y\n\nresult = f(2.)\n```\n\nExample:\n```text\njax.debug.print(x) -> 2.0\njax.debug.print(y) -> 0.9092974066734314\n```\n\nExample:\n```text\ndef f(x):\n  jax.debug.print(\"jax.debug.print(x) -> {}\", x)\n  y = jnp.sin(x)\n  jax.debug.print(\"jax.debug.print(y) -> {}\", y)\n  return y\n\nxs = jnp.arange(3.)\n\nresult = jax.vmap(f)(xs)\n```\n\nExample:\n```text\njax.debug.print(x) -> 0.0\njax.debug.print(x) -> 1.0\njax.debug.print(x) -> 2.0\njax.debug.print(y) -> 0.0\njax.debug.print(y) -> 0.8414709568023682\njax.debug.print(y) -> 0.9092974066734314\n```\n\nExample:\n```text\nresult = jax.lax.map(f, xs)\n```\n\nExample:\n```text\njax.debug.print(y) -> 0.0\njax.debug.print(x) -> 0.0\njax.debug.print(y) -> 0.8414709568023682\njax.debug.print(x) -> 1.0\njax.debug.print(y) -> 0.9092974066734314\njax.debug.print(x) -> 2.0\n```\n\nExample:\n```text\ndef f(x):\n  jax.debug.print(\"jax.debug.print(x) -> {}\", x)\n  return x ** 2\n\nresult = jax.grad(f)(1.)\n```\n\nExample:\n```text\njax.debug.print(x) -> 1.0\n```\n\nExample:\n```text\n@jax.jit\ndef f(x, y):\n  jax.debug.print(\"jax.debug.print(x) -> {}\", x, ordered=True)\n  jax.debug.print(\"jax.debug.print(y) -> {}\", y, ordered=True)\n  return x + y\n\nf(1, 2)\n```\n\nExample:\n```text\njax.debug.print(x) -> 1\njax.debug.print(y) -> 2\n```\n\nExample:\n```text\nArray(3, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  y, z = jnp.sin(x), jnp.cos(x)\n  jax.debug.breakpoint()\n  return y * z\n\nf(2.) # ==> Pauses during execution\n```\n\nExample:\n```text\ndef breakpoint_if_nonfinite(x):\n  is_finite = jnp.isfinite(x).all()\n  def true_fn(x):\n    pass\n  def false_fn(x):\n    jax.debug.breakpoint()\n  jax.lax.cond(is_finite, true_fn, false_fn, x)\n\n@jax.jit\ndef f(x, y):\n  z = x / y\n  breakpoint_if_nonfinite(z)\n  return z\n\nf(2., 1.) # ==> No breakpoint\n```\n\nExample:\n```text\nArray(2., dtype=float32, weak_type=True)\n```\n\nExample:\n```text\nf(2., 0.) # ==> Pauses during execution\n```\n\nExample:\n```text\nimport logging\n\ndef log_value(x):\n  logging.warning(f'Logged value: {x}')\n\n@jax.jit\ndef f(x):\n  jax.debug.callback(log_value, x)\n  return x\n\nf(1.0);\n```\n\nExample:\n```text\nWARNING:root:Logged value: 1.0\n```\n\nExample:\n```text\nx = jnp.arange(5.0)\njax.vmap(f)(x);\n```\n\nExample:\n```text\nWARNING:root:Logged value: 0.0\nWARNING:root:Logged value: 1.0\nWARNING:root:Logged value: 2.0\nWARNING:root:Logged value: 3.0\nWARNING:root:Logged value: 4.0\n```\n\nExample:\n```text\njax.grad(f)(1.0);\n```\n\nExample:\n```text\nfrom jax.experimental import checkify\nimport jax\nimport jax.numpy as jnp\n\ndef f(x, i):\n  checkify.check(i >= 0, \"index needs to be non-negative!\")\n  y = x[i]\n  z = jnp.sin(y)\n  return z\n\njittable_f = checkify.checkify(f)\n\nerr, z = jax.jit(jittable_f)(jnp.ones((5,)), -1)\nprint(err.get())\n# >> index needs to be non-negative! (check failed at <...>:6 (f))\n```\n\nExample:\n```text\nerrors = checkify.user_checks | checkify.index_checks | checkify.float_checks\nchecked_f = checkify.checkify(f, errors=errors)\n\nerr, z = checked_f(jnp.ones((5,)), 100)\nerr.throw()\n# ValueError: out-of-bounds indexing at <..>:7 (f)\n\nerr, z = checked_f(jnp.ones((5,)), -1)\nerr.throw()\n# ValueError: index needs to be non-negative! (check failed at <…>:6 (f))\n\nerr, z = checked_f(jnp.array([jnp.inf, 1]), 0)\nerr.throw()\n# ValueError: nan generated by primitive sin at <...>:8 (f)\n```\n\nExample:\n```text\nimport jax\njax.config.update(\"jax_debug_nans\", True)\n\ndef f(x, y):\n  return x / y\n\njax.jit(f)(0., 0.)  # ==> raises FloatingPointError exception!\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.681Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":25,"totalLines":242,"estimatedTokens":964}}18{"id":"doc-forward_and_reverse_mode_autodiff_in_jax_jax_doc-0d1c7fff","source":"documentation","title":"Forward- and reverse-mode autodiff in JAX — JAX documentation","url":"https://docs.jax.dev/en/latest/jacobian-vector-products.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n\nkey = jax.random.key(0)\n\n# Initialize random model coefficients\nkey, W_key, b_key = jax.random.split(key, 3)\nW = jax.random.normal(W_key, (3,))\nb = jax.random.normal(b_key, ())\n\n# Define a sigmoid function.\ndef sigmoid(x):\n    return 0.5 * (jnp.tanh(x / 2) + 1)\n\n# Outputs probability of a label being true.\ndef predict(W, b, inputs):\n    return sigmoid(jnp.dot(inputs, W) + b)\n\n# Build a toy dataset.\ninputs = jnp.array([[0.52, 1.12,  0.77],\n                   [0.88, -1.08, 0.15],\n                   [0.52, 0.06, -1.30],\n                   [0.74, -2.49, 1.39]])\n\n# Isolate the function from the weight matrix to the predictions\nf = lambda W: predict(W, b, inputs)\n\nkey, subkey = jax.random.split(key)\nv = jax.random.normal(subkey, W.shape)\n\n# Push forward the vector `v` along `f` evaluated at `W`\ny, u = jax.jvp(f, (W,), (v,))\n```\n\nExample:\n```text\njvp :: (a -> b) -> a -> T a -> (b, T b)\n```\n\nExample:\n```text\nfrom jax import vjp\n\n# Isolate the function from the weight matrix to the predictions\nf = lambda W: predict(W, b, inputs)\n\ny, vjp_fun = vjp(f, W)\n\nkey, subkey = jax.random.split(key)\nu = jax.random.normal(subkey, y.shape)\n\n# Pull back the covector `u` along `f` evaluated at `W`\nv = vjp_fun(u)\n```\n\nExample:\n```text\nvjp :: (a -> b) -> a -> (b, CT b -> CT a)\n```\n\nExample:\n```text\ndef vgrad(f, x):\n  y, vjp_fn = jax.vjp(f, x)\n  return vjp_fn(jnp.ones(y.shape))[0]\n\nprint(vgrad(lambda x: 3*x**2, jnp.ones((2, 2))))\n```\n\nExample:\n```text\n[[6. 6.]\n [6. 6.]]\n```\n\nExample:\n```text\ndef hvp(f, x, v):\n    return jax.grad(lambda x: jnp.vdot(jax.grad(f)(x), v))(x)\n```\n\nExample:\n```text\n# forward-over-reverse\ndef hvp(f, primals, tangents):\n  return jax.jvp(jax.grad(f), primals, tangents)[1]\n```\n\nExample:\n```text\ndef f(X):\n  return jnp.sum(jnp.tanh(X)**2)\n\nkey, subkey1, subkey2 = jax.random.split(key, 3)\nX = jax.random.normal(subkey1, (30, 40))\nV = jax.random.normal(subkey2, (30, 40))\n\ndef hessian(f):\n    return jax.jacfwd(jax.jacrev(f))\n\nans1 = hvp(f, (X,), (V,))\nans2 = jnp.tensordot(hessian(f)(X), V, 2)\n\nprint(jnp.allclose(ans1, ans2, 1e-4, 1e-4))\n```\n\nExample:\n```text\nTrue\n```\n\nExample:\n```text\n# Reverse-over-forward\ndef hvp_revfwd(f, primals, tangents):\n  g = lambda primals: jax.jvp(f, primals, tangents)[1]\n  return jax.grad(g)(primals)\n```\n\nExample:\n```text\n# Reverse-over-reverse, only works for single arguments\ndef hvp_revrev(f, primals, tangents):\n  x, = primals\n  v, = tangents\n  return jax.grad(lambda x: jnp.vdot(jax.grad(f)(x), v))(x)\n\n\nprint(\"Forward over reverse\")\n%timeit -n10 -r3 hvp(f, (X,), (V,))\nprint(\"Reverse over forward\")\n%timeit -n10 -r3 hvp_revfwd(f, (X,), (V,))\nprint(\"Reverse over reverse\")\n%timeit -n10 -r3 hvp_revrev(f, (X,), (V,))\n\nprint(\"Naive full Hessian materialization\")\n%timeit -n10 -r3 jnp.tensordot(jax.hessian(f)(X), V, 2)\n```\n\nExample:\n```text\nForward over reverse\n2.66 ms ± 92.7 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)\nReverse over forward\nThe slowest run took 4.75 times longer than the fastest. This could mean that an intermediate result is being cached.\n8.83 ms ± 6.94 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)\nReverse over reverse\n11.3 ms ± 7.11 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)\nNaive full Hessian materialization\n40.8 ms ± 1.36 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)\n```\n\nExample:\n```text\n# Isolate the function from the weight matrix to the predictions\nf = lambda W: predict(W, b, inputs)\n\n# Pull back the covectors `m_i` along `f`, evaluated at `W`, for all `i`.\n# First, use a list comprehension to loop over rows in the matrix M.\ndef loop_mjp(f, x, M):\n    y, vjp_fun = jax.vjp(f, x)\n    return jnp.vstack([jnp.asarray(vjp_fun(mi)) for mi in M])\n\n# Now, use vmap to build a computation that does a single fast matrix-matrix\n# multiply, rather than an outer loop over vector-matrix multiplies.\ndef vmap_mjp(f, x, M):\n    y, vjp_fun = jax.vjp(f, x)\n    outs, = jax.vmap(vjp_fun)(M)\n    return outs\n\nkey = jax.random.key(0)\nnum_covecs = 128\nU = jax.random.normal(key, (num_covecs,) + y.shape)\n\nloop_vs = loop_mjp(f, W, M=U)\nprint('Non-vmapped Matrix-Jacobian product')\n%timeit -n10 -r3 loop_mjp(f, W, M=U)\n\nprint('\\nVmapped Matrix-Jacobian product')\nvmap_vs = vmap_mjp(f, W, M=U)\n%timeit -n10 -r3 vmap_mjp(f, W, M=U)\n\nassert jnp.allclose(loop_vs, vmap_vs), 'Vmap and non-vmapped Matrix-Jacobian Products should be identical'\n```\n\nExample:\n```text\nNon-vmapped Matrix-Jacobian product\n59.6 ms ± 160 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)\n\nVmapped Matrix-Jacobian product\n3.39 ms ± 43.7 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)\n```\n\nExample:\n```text\ndef loop_jmp(f, W, M):\n    # jvp immediately returns the primal and tangent values as a tuple,\n    # so we'll compute and select the tangents in a list comprehension\n    return jnp.vstack([jax.jvp(f, (W,), (mi,))[1] for mi in M])\n\ndef vmap_jmp(f, W, M):\n    _jvp = lambda s: jax.jvp(f, (W,), (s,))[1]\n    return jax.vmap(_jvp)(M)\nnum_vecs = 128\nS = jax.random.normal(key, (num_vecs,) + W.shape)\n\nloop_vs = loop_jmp(f, W, M=S)\nprint('Non-vmapped Jacobian-Matrix product')\n%timeit -n10 -r3 loop_jmp(f, W, M=S)\nvmap_vs = vmap_jmp(f, W, M=S)\nprint('\\nVmapped Jacobian-Matrix product')\n%timeit -n10 -r3 vmap_jmp(f, W, M=S)\n\nassert jnp.allclose(loop_vs, vmap_vs), 'Vmap and non-vmapped Jacobian-Matrix products should be identical'\n```\n\nExample:\n```text\nNon-vmapped Jacobian-Matrix product\n78.3 ms ± 232 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)\n\nVmapped Jacobian-Matrix product\n1.18 ms ± 36.2 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)\n```\n\nExample:\n```text\nfrom jax import jacrev as builtin_jacrev\n\ndef our_jacrev(f):\n    def jacfun(x):\n        y, vjp_fun = jax.vjp(f, x)\n        # Use vmap to do a matrix-Jacobian product.\n        # Here, the matrix is the Euclidean basis, so we get all\n        # entries in the Jacobian at once.\n        J, = jax.vmap(vjp_fun, in_axes=0)(jnp.eye(len(y)))\n        return J\n    return jacfun\n\nassert jnp.allclose(builtin_jacrev(f)(W), our_jacrev(f)(W)), 'Incorrect reverse-mode Jacobian results!'\n```\n\nExample:\n```text\nfrom jax import jacfwd as builtin_jacfwd\n\ndef our_jacfwd(f):\n    def jacfun(x):\n        _jvp = lambda s: jax.jvp(f, (x,), (s,))[1]\n        Jt = jax.vmap(_jvp, in_axes=1)(jnp.eye(len(x)))\n        return jnp.transpose(Jt)\n    return jacfun\n\nassert jnp.allclose(builtin_jacfwd(f)(W), our_jacfwd(f)(W)), 'Incorrect forward-mode Jacobian results!'\n```\n\nExample:\n```text\ndef f(x):\n    try:\n        if x < 3:\n            return 2 * x ** 3\n        else:\n            raise ValueError\n    except ValueError:\n        return jnp.pi * x\n\ny, f_vjp = jax.vjp(f, 4.)\nprint(jax.jit(f_vjp)(1.))\n```\n\nExample:\n```text\n(Array(3.1415927, dtype=float32, weak_type=True),)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.682Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":21,"totalLines":280,"estimatedTokens":1703}}19{"id":"doc-jax_numpy_fft_fftshift_jax_documentation-6bd88b7f","source":"documentation","title":"jax.numpy.fft.fftshift — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.fftshift.html","text":"Example:\n```text\n>>> freq = jnp.fft.fftfreq(5)\n>>> freq\nArray([ 0. ,  0.2,  0.4, -0.4, -0.2], dtype=float32)\n```\n\nExample:\n```text\n>>> shifted_freq = jnp.fft.fftshift(freq)\n>>> shifted_freq\nArray([-0.4, -0.2,  0. ,  0.2,  0.4], dtype=float32)\n```\n\nExample:\n```text\n>>> jnp.fft.ifftshift(shifted_freq)\nArray([ 0. ,  0.2,  0.4, -0.4, -0.2], dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.682Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":21,"estimatedTokens":93}}20{"id":"doc-jax_numpy_fft_irfft_jax_documentation-eede3696","source":"documentation","title":"jax.numpy.fft.irfft — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.irfft.html","text":"Example:\n```text\n>>> x = jnp.array([[1, 3, 5],\n...                [2, 4, 6]])\n>>> jnp.fft.irfft(x)\nArray([[ 3., -1.,  0., -1.],\n       [ 4., -1.,  0., -1.]], dtype=float32)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.irfft(x, n=3)\nArray([[ 2.33, -0.67, -0.67],\n       [ 3.33, -0.67, -0.67]], dtype=float32)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.irfft(x, n=4, axis=0)\nArray([[ 1.25,  2.75,  4.25],\n       [ 0.25,  0.75,  1.25],\n       [-0.75, -1.25, -1.75],\n       [ 0.25,  0.75,  1.25]], dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.683Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":28,"estimatedTokens":155}}21{"id":"doc-higher_order_derivatives_jax_documentation-53efd2b2","source":"documentation","title":"Higher-order derivatives — JAX documentation","url":"https://docs.jax.dev/en/latest/higher-order.html","text":"Example:\n```text\nimport jax\n\ndef hessian(f):\n  return jax.jacfwd(jax.grad(f))\n```\n\nExample:\n```text\nimport jax.numpy as jnp\n\ndef f(x):\n  return jnp.dot(x, x)\n\nhessian(f)(jnp.array([1., 2., 3.]))\n```\n\nExample:\n```text\nArray([[2., 0., 0.],\n       [0., 2., 0.],\n       [0., 0., 2.]], dtype=float32)\n```\n\nExample:\n```text\ndef meta_loss_fn(params, data):\n  \"\"\"Computes the loss after one step of SGD.\"\"\"\n  grads = jax.grad(loss_fn)(params, data)\n  return loss_fn(params - lr * grads, data)\n\nmeta_grads = jax.grad(meta_loss_fn)(params, data)\n```\n\nExample:\n```text\n# Value function and initial parameters\nvalue_fn = lambda theta, state: jnp.dot(theta, state)\ntheta = jnp.array([0.1, -0.1, 0.])\n```\n\nExample:\n```text\n# An example transition.\ns_tm1 = jnp.array([1., 2., -1.])\nr_t = jnp.array(1.)\ns_t = jnp.array([2., 1., 0.])\n```\n\nExample:\n```text\ndef td_loss(theta, s_tm1, r_t, s_t):\n  v_tm1 = value_fn(theta, s_tm1)\n  target = r_t + value_fn(theta, s_t)\n  return -0.5 * ((target - v_tm1) ** 2)\n\ntd_update = jax.grad(td_loss)\ndelta_theta = td_update(theta, s_tm1, r_t, s_t)\n\ndelta_theta\n```\n\nExample:\n```text\nArray([-1.2,  1.2, -1.2], dtype=float32)\n```\n\nExample:\n```text\ndef td_loss(theta, s_tm1, r_t, s_t):\n  v_tm1 = value_fn(theta, s_tm1)\n  target = r_t + value_fn(theta, s_t)\n  return -0.5 * ((jax.lax.stop_gradient(target) - v_tm1) ** 2)\n\ntd_update = jax.grad(td_loss)\ndelta_theta = td_update(theta, s_tm1, r_t, s_t)\n\ndelta_theta\n```\n\nExample:\n```text\nArray([ 1.2,  2.4, -1.2], dtype=float32)\n```\n\nExample:\n```text\ns_grad = jax.grad(value_fn)(theta, s_tm1)\ndelta_theta_original_calculation = (r_t + value_fn(theta, s_t) - value_fn(theta, s_tm1)) * s_grad\n\ndelta_theta_original_calculation # [1.2, 2.4, -1.2], same as `delta_theta`\n```\n\nExample:\n```text\ndef f(x):\n  return jnp.round(x)  # non-differentiable\n\ndef straight_through_f(x):\n  # Create an exactly-zero expression with Sterbenz lemma that has\n  # an exactly-one gradient.\n  zero = x - jax.lax.stop_gradient(x)\n  return zero + jax.lax.stop_gradient(f(x))\n\nprint(\"f(x): \", f(3.2))\nprint(\"straight_through_f(x):\", straight_through_f(3.2))\n\nprint(\"grad(f)(x):\", jax.grad(f)(3.2))\nprint(\"grad(straight_through_f)(x):\", jax.grad(straight_through_f)(3.2))\n```\n\nExample:\n```text\nf(x):  3.0\nstraight_through_f(x): 3.0\ngrad(f)(x): 0.0\ngrad(straight_through_f)(x): 1.0\n```\n\nExample:\n```text\nperex_grads = jax.jit(jax.vmap(jax.grad(td_loss), in_axes=(None, 0, 0, 0)))\n\n# Test it:\nbatched_s_tm1 = jnp.stack([s_tm1, s_tm1])\nbatched_r_t = jnp.stack([r_t, r_t])\nbatched_s_t = jnp.stack([s_t, s_t])\n\nperex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t)\n```\n\nExample:\n```text\nArray([[ 1.2,  2.4, -1.2],\n       [ 1.2,  2.4, -1.2]], dtype=float32)\n```\n\nExample:\n```text\ndtdloss_dtheta = jax.grad(td_loss)\n\ndtdloss_dtheta(theta, s_tm1, r_t, s_t)\n```\n\nExample:\n```text\nalmost_perex_grads = jax.vmap(dtdloss_dtheta)\n\nbatched_theta = jnp.stack([theta, theta])\nalmost_perex_grads(batched_theta, batched_s_tm1, batched_r_t, batched_s_t)\n```\n\nExample:\n```text\ninefficient_perex_grads = jax.vmap(dtdloss_dtheta, in_axes=(None, 0, 0, 0))\n\ninefficient_perex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t)\n```\n\nExample:\n```text\nperex_grads = jax.jit(inefficient_perex_grads)\n\nperex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t)\n```\n\nExample:\n```text\n%timeit inefficient_perex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t).block_until_ready()\n%timeit perex_grads(theta, batched_s_tm1, batched_r_t, batched_s_t).block_until_ready()\n```\n\nExample:\n```text\n3.03 ms ± 6.2 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n6.79 μs ± 25.1 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n```\n\nExample:\n```text\ndef hvp(f, x, v):\n    return jax.grad(lambda x: jnp.vdot(jax.grad(f)(x), v))(x)\n```\n\nExample:\n```text\nfrom jax import jacfwd, jacrev\n\n# Define a sigmoid function.\ndef sigmoid(x):\n    return 0.5 * (jnp.tanh(x / 2) + 1)\n\n# Outputs probability of a label being true.\ndef predict(W, b, inputs):\n    return sigmoid(jnp.dot(inputs, W) + b)\n\n# Build a toy dataset.\ninputs = jnp.array([[0.52, 1.12,  0.77],\n                   [0.88, -1.08, 0.15],\n                   [0.52, 0.06, -1.30],\n                   [0.74, -2.49, 1.39]])\n\n# Initialize random model coefficients\nkey = jax.random.key(0)\nkey, W_key, b_key = jax.random.split(key, 3)\nW = jax.random.normal(W_key, (3,))\nb = jax.random.normal(b_key, ())\n\n# Isolate the function from the weight matrix to the predictions\nf = lambda W: predict(W, b, inputs)\n\nJ = jacfwd(f)(W)\nprint(\"jacfwd result, with shape\", J.shape)\nprint(J)\n\nJ = jacrev(f)(W)\nprint(\"jacrev result, with shape\", J.shape)\nprint(J)\n```\n\nExample:\n```text\njacfwd result, with shape (4, 3)\n[[ 0.05069415  0.1091874   0.07506633]\n [ 0.14170025 -0.17390487  0.02415345]\n [ 0.12579198  0.01451446 -0.31447992]\n [ 0.00574409 -0.0193281   0.01078958]]\njacrev result, with shape (4, 3)\n[[ 0.05069415  0.10918741  0.07506634]\n [ 0.14170025 -0.17390487  0.02415345]\n [ 0.12579198  0.01451446 -0.31447995]\n [ 0.00574409 -0.0193281   0.01078958]]\n```\n\nExample:\n```text\ndef predict_dict(params, inputs):\n    return predict(params['W'], params['b'], inputs)\n\nJ_dict = jax.jacrev(predict_dict)({'W': W, 'b': b}, inputs)\nfor k, v in J_dict.items():\n    print(\"Jacobian from {} to logits is\".format(k))\n    print(v)\n```\n\nExample:\n```text\nJacobian from W to logits is\n[[ 0.05069415  0.10918741  0.07506634]\n [ 0.14170025 -0.17390487  0.02415345]\n [ 0.12579198  0.01451446 -0.31447995]\n [ 0.00574409 -0.0193281   0.01078958]]\nJacobian from b to logits is\n[0.09748876 0.16102302 0.24190766 0.00776229]\n```\n\nExample:\n```text\ndef hessian(f):\n    return jax.jacfwd(jax.jacrev(f))\n\nH = hessian(f)(W)\nprint(\"hessian, with shape\", H.shape)\nprint(H)\n```\n\nExample:\n```text\nhessian, with shape (4, 3, 3)\n[[[ 0.02058932  0.04434624  0.03048803]\n  [ 0.04434623  0.09551499  0.06566654]\n  [ 0.03048803  0.06566655  0.04514575]]\n\n [[-0.0743913   0.09129842 -0.01268033]\n  [ 0.09129842 -0.11204806  0.01556223]\n  [-0.01268034  0.01556223 -0.00216142]]\n\n [[ 0.01176856  0.00135791 -0.02942139]\n  [ 0.00135791  0.00015668 -0.00339478]\n  [-0.0294214  -0.00339478  0.07355348]]\n\n [[-0.00418412  0.014079   -0.00785936]\n  [ 0.014079   -0.04737393  0.02644569]\n  [-0.00785936  0.02644569 -0.01476286]]]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.684Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":28,"totalLines":288,"estimatedTokens":1566}}22{"id":"doc-jax_numpy_fft_ifftn_jax_documentation-9f6c68e9","source":"documentation","title":"jax.numpy.fft.ifftn — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.ifftn.html","text":"Example:\n```text\n>>> x = jnp.array([[1, 2, 5, 3],\n...                [4, 1, 2, 6],\n...                [5, 3, 2, 1]])\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jnp.fft.ifftn(x))\n[[ 2.92+0.j    0.08-0.33j  0.25+0.j    0.08+0.33j]\n [-0.08+0.14j -0.04-0.03j  0.  -0.29j -1.05-0.11j]\n [-0.08-0.14j -1.05+0.11j  0.  +0.29j -0.04+0.03j]]\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jnp.fft.ifftn(x, s=[3]))\n[[ 2.67+0.j   -0.83-0.87j -0.83+0.87j]\n [ 2.33+0.j    0.83-0.29j  0.83+0.29j]\n [ 3.33+0.j    0.83+0.29j  0.83-0.29j]]\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jnp.fft.ifftn(x, s=[2], axes=[0]))\n[[ 2.5+0.j  1.5+0.j  3.5+0.j  4.5+0.j]\n [-1.5+0.j  0.5+0.j  1.5+0.j -1.5+0.j]]\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jnp.fft.ifftn(x, s=[2, 3]))\n[[ 2.5 +0.j    0.  -0.58j  0.  +0.58j]\n [ 0.17+0.j   -0.83-0.29j -0.83+0.29j]]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.684Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":38,"estimatedTokens":251}}23{"id":"doc-jax_numpy_fft_ifftshift_jax_documentation-cbfe5b88","source":"documentation","title":"jax.numpy.fft.ifftshift — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.ifftshift.html","text":"Example:\n```text\n>>> freq = jnp.fft.fftfreq(5)\n>>> freq\nArray([ 0. ,  0.2,  0.4, -0.4, -0.2], dtype=float32)\n```\n\nExample:\n```text\n>>> shifted_freq = jnp.fft.fftshift(freq)\n>>> shifted_freq\nArray([-0.4, -0.2,  0. ,  0.2,  0.4], dtype=float32)\n```\n\nExample:\n```text\n>>> jnp.fft.ifftshift(shifted_freq)\nArray([ 0. ,  0.2,  0.4, -0.4, -0.2], dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.685Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":21,"estimatedTokens":93}}24{"id":"doc-jax_numpy_fft_fft_jax_documentation-819a4d80","source":"documentation","title":"jax.numpy.fft.fft — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.fft.html","text":"Example:\n```text\n>>> x = jnp.array([[1, 2, 4, 7],\n...                [5, 3, 1, 9]])\n>>> jnp.fft.fft(x)\nArray([[14.+0.j, -3.+5.j, -4.+0.j, -3.-5.j],\n       [18.+0.j,  4.+6.j, -6.+0.j,  4.-6.j]], dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jnp.fft.fft(x, n=3))\n[[ 7.+0.j   -2.+1.73j -2.-1.73j]\n [ 9.+0.j    3.-1.73j  3.+1.73j]]\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jnp.fft.fft(x, n=3, axis=0))\n[[ 6. +0.j    5. +0.j    5. +0.j   16. +0.j  ]\n [-1.5-4.33j  0.5-2.6j   3.5-0.87j  2.5-7.79j]\n [-1.5+4.33j  0.5+2.6j   3.5+0.87j  2.5+7.79j]]\n```\n\nExample:\n```text\n>>> x_fft = jnp.fft.fft(x)\n>>> jnp.allclose(x, jnp.fft.ifft(x_fft))\nArray(True, dtype=bool)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.685Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":34,"estimatedTokens":195}}25{"id":"doc-jax_numpy_fft_ihfft_jax_documentation-2ab165be","source":"documentation","title":"jax.numpy.fft.ihfft — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.ihfft.html","text":"Example:\n```text\n>>> x = jnp.array([[1, 3, 5, 7],\n...                [2, 4, 6, 8]])\n>>> jnp.fft.ihfft(x)\nArray([[ 4.+0.j, -1.-1.j, -1.-0.j],\n       [ 5.+0.j, -1.-1.j, -1.-0.j]], dtype=complex64)\n```\n\nExample:\n```text\n>>> jnp.fft.ihfft(x, n=4, axis=0)\nArray([[ 0.75+0.j ,  1.75+0.j ,  2.75+0.j ,  3.75+0.j ],\n       [ 0.25+0.5j,  0.75+1.j ,  1.25+1.5j,  1.75+2.j ],\n       [-0.25-0.j , -0.25-0.j , -0.25-0.j , -0.25-0.j ]], dtype=complex64)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.685Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":2,"totalLines":18,"estimatedTokens":115}}26{"id":"doc-jax_numpy_fft_fftn_jax_documentation-3c7cf928","source":"documentation","title":"jax.numpy.fft.fftn — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.fftn.html","text":"Example:\n```text\n>>> x = jnp.array([[1, 2, 5, 6],\n...                [4, 1, 3, 7],\n...                [5, 9, 2, 1]])\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.fftn(x)\nArray([[ 46.  +0.j  ,   0.  +2.j  ,  -6.  +0.j  ,   0.  -2.j  ],\n       [ -2.  +1.73j,   6.12+6.73j,   0.  -1.73j, -18.12-3.27j],\n       [ -2.  -1.73j, -18.12+3.27j,   0.  +1.73j,   6.12-6.73j]],      dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jax.numpy.fft.fftn(x, s=[2]))\n[[ 3.+0.j -1.+0.j]\n [ 5.+0.j  3.+0.j]\n [14.+0.j -4.+0.j]]\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jax.numpy.fft.fftn(x, s=[2], axes=[0]))\n[[ 5.+0.j  3.+0.j  8.+0.j 13.+0.j]\n [-3.+0.j  1.+0.j  2.+0.j -1.+0.j]]\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jax.numpy.fft.fftn(x, s=[2, 3]))\n[[16. +0.j   -0.5+4.33j -0.5-4.33j]\n [ 0. +0.j   -4.5+0.87j -4.5-0.87j]]\n```\n\nExample:\n```text\n>>> x_fftn = jnp.fft.fftn(x)\n>>> jnp.allclose(x, jnp.fft.ifftn(x_fftn))\nArray(True, dtype=bool)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.686Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":5,"totalLines":45,"estimatedTokens":280}}27{"id":"doc-jax_numpy_fft_ifft_jax_documentation-6448957d","source":"documentation","title":"jax.numpy.fft.ifft — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.ifft.html","text":"Example:\n```text\n>>> x = jnp.array([[3, 1, 4, 6],\n...                [2, 5, 7, 1]])\n>>> jnp.fft.ifft(x)\nArray([[ 3.5 +0.j  , -0.25-1.25j,  0.  +0.j  , -0.25+1.25j],\n      [ 3.75+0.j  , -1.25+1.j  ,  0.75+0.j  , -1.25-1.j  ]],      dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jnp.fft.ifft(x, n=5))\n[[ 2.8 +0.j   -0.96-0.04j  1.06+0.5j   1.06-0.5j  -0.96+0.04j]\n [ 3.  +0.j   -0.59+1.66j  0.09-0.55j  0.09+0.55j -0.59-1.66j]]\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jnp.fft.ifft(x, n=3, axis=0))\n[[ 1.67+0.j    2.  +0.j    3.67+0.j    2.33+0.j  ]\n [ 0.67+0.58j -0.5 +1.44j  0.17+2.02j  1.83+0.29j]\n [ 0.67-0.58j -0.5 -1.44j  0.17-2.02j  1.83-0.29j]]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.686Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":27,"estimatedTokens":194}}28{"id":"doc-jax_numpy_fft_irfft2_jax_documentation-a54188c1","source":"documentation","title":"jax.numpy.fft.irfft2 — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.irfft2.html","text":"Example:\n```text\n>>> x = jnp.array([[[1, 3, 5],\n...                 [2, 4, 6]],\n...                [[7, 9, 11],\n...                 [8, 10, 12]]])\n>>> jnp.fft.irfft2(x)\nArray([[[ 3.5, -1. ,  0. , -1. ],\n        [-0.5,  0. ,  0. ,  0. ]],\n\n       [[ 9.5, -1. ,  0. , -1. ],\n        [-0.5,  0. ,  0. ,  0. ]]], dtype=float32)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.irfft2(x, s=[3, 3])\nArray([[[ 1.89, -0.44, -0.44],\n        [ 0.22, -0.78,  0.56],\n        [ 0.22,  0.56, -0.78]],\n\n       [[ 5.89, -0.44, -0.44],\n        [ 1.22, -1.78,  1.56],\n        [ 1.22,  1.56, -1.78]]], dtype=float32)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.irfft2(x, s=[2, 3], axes=(0, 1))\nArray([[[ 4.67,  6.67,  8.67],\n        [-0.33, -0.33, -0.33],\n        [-0.33, -0.33, -0.33]],\n\n       [[-3.  , -3.  , -3.  ],\n        [ 0.  ,  0.  ,  0.  ],\n        [ 0.  ,  0.  ,  0.  ]]], dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.686Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":41,"estimatedTokens":247}}29{"id":"doc-jax_numpy_fft_fft2_jax_documentation-c89f72ef","source":"documentation","title":"jax.numpy.fft.fft2 — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.fft2.html","text":"Example:\n```text\n>>> x = jnp.array([[[1, 3],\n...                 [2, 4]],\n...                [[5, 7],\n...                 [6, 8]]])\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.fft2(x)\nArray([[[10.+0.j, -4.+0.j],\n        [-2.+0.j,  0.+0.j]],\n\n       [[26.+0.j, -4.+0.j],\n        [-2.+0.j,  0.+0.j]]], dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.fft2(x, s=[2, 3])\nArray([[[10.  +0.j  , -0.5 -6.06j, -0.5 +6.06j],\n        [-2.  +0.j  , -0.5 +0.87j, -0.5 -0.87j]],\n\n       [[26.  +0.j  ,  3.5-12.99j,  3.5+12.99j],\n        [-2.  +0.j  , -0.5 +0.87j, -0.5 -0.87j]]], dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.fft2(x, s=[2, 3], axes=(0, 1))\nArray([[[14. +0.j  , 22. +0.j  ],\n        [ 2. -6.93j,  4.-10.39j],\n        [ 2. +6.93j,  4.+10.39j]],\n\n       [[-8. +0.j  , -8. +0.j  ],\n        [-2. +3.46j, -2. +3.46j],\n        [-2. -3.46j, -2. -3.46j]]], dtype=complex64)\n```\n\nExample:\n```text\n>>> x_fft2 = jnp.fft.fft2(x)\n>>> jnp.allclose(x, jnp.fft.ifft2(x_fft2))\nArray(True, dtype=bool)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.687Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":47,"estimatedTokens":288}}30{"id":"doc-jax_numpy_fft_hfft_jax_documentation-4e84a8cf","source":"documentation","title":"jax.numpy.fft.hfft — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.hfft.html","text":"Example:\n```text\n>>> x = jnp.array([[1, 3, 5, 7],\n...                [2, 4, 6, 8]])\n>>> jnp.fft.hfft(x)\nArray([[24., -8.,  0., -2.,  0., -8.],\n       [30., -8.,  0., -2.,  0., -8.]], dtype=float32)\n```\n\nExample:\n```text\n>>> x1 = jnp.array([[1, 3, 5, 7, 5, 3],\n...                 [2, 4, 6, 8, 6, 4]])\n>>> jnp.fft.fft(x1)\nArray([[24.+0.j, -8.+0.j,  0.+0.j, -2.+0.j,  0.+0.j, -8.+0.j],\n       [30.+0.j, -8.+0.j,  0.+0.j, -2.+0.j,  0.+0.j, -8.+0.j]],      dtype=complex64)\n>>> jnp.allclose(jnp.fft.hfft(x), jnp.fft.fft(x1))\nArray(True, dtype=bool)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   print(jnp.fft.hfft(x, n=5))\n[[17.   -5.24 -0.76 -0.76 -5.24]\n [22.   -5.24 -0.76 -0.76 -5.24]]\n```\n\nExample:\n```text\n>>> jnp.fft.hfft(x, n=3, axis=0)\nArray([[ 5., 11., 17., 23.],\n       [-1., -1., -1., -1.],\n       [-1., -1., -1., -1.]], dtype=float32)\n```\n\nExample:\n```text\n>>> jnp.fft.ihfft(jnp.fft.hfft(x, 2*(x.shape[-1]-1)))\nArray([[1.+0.j, 3.+0.j, 5.+0.j, 7.+0.j],\n       [2.+0.j, 4.+0.j, 6.+0.j, 8.+0.j]], dtype=complex64)\n>>> jnp.allclose(x, jnp.fft.ihfft(jnp.fft.hfft(x, 2*(x.shape[-1]-1))))\nArray(True, dtype=bool)\n```\n\nExample:\n```text\n>>> x2 = jnp.array([[1+2j, 3-4j, 5+6j],\n...                 [2-3j, 4+5j, 6-7j]])\n>>> jnp.fft.hfft(x2)\nArray([[ 12., -12.,   0.,   4.],\n       [ 16.,   6.,   0., -14.]], dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.687Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":6,"totalLines":55,"estimatedTokens":344}}31{"id":"doc-jax_numpy_fft_irfftn_jax_documentation-21b4a171","source":"documentation","title":"jax.numpy.fft.irfftn — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.irfftn.html","text":"Example:\n```text\n>>> x = jnp.array([[[1, 3, 5],\n...                 [2, 4, 6]],\n...                [[7, 9, 11],\n...                 [8, 10, 12]]])\n>>> jnp.fft.irfftn(x)\nArray([[[ 6.5, -1. ,  0. , -1. ],\n        [-0.5,  0. ,  0. ,  0. ]],\n\n       [[-3. ,  0. ,  0. ,  0. ],\n        [ 0. ,  0. ,  0. ,  0. ]]], dtype=float32)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.irfftn(x, s=[3, 4])\nArray([[[ 2.33, -0.67,  0.  , -0.67],\n        [ 0.33, -0.74,  0.  ,  0.41],\n        [ 0.33,  0.41,  0.  , -0.74]],\n\n       [[ 6.33, -0.67,  0.  , -0.67],\n        [ 1.33, -1.61,  0.  ,  1.28],\n        [ 1.33,  1.28,  0.  , -1.61]]], dtype=float32)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.irfftn(x, s=[3], axes=[0])\nArray([[[ 5.,  7.,  9.],\n        [ 6.,  8., 10.]],\n\n       [[-2., -2., -2.],\n        [-2., -2., -2.]],\n\n       [[-2., -2., -2.],\n        [-2., -2., -2.]]], dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.688Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":42,"estimatedTokens":247}}32{"id":"doc-jax_numpy_fft_ifft2_jax_documentation-6575c230","source":"documentation","title":"jax.numpy.fft.ifft2 — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.ifft2.html","text":"Example:\n```text\n>>> x = jnp.array([[[1, 3],\n...                 [2, 4]],\n...                [[5, 7],\n...                 [6, 8]]])\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.ifft2(x)\nArray([[[ 2.5+0.j, -1. +0.j],\n        [-0.5+0.j,  0. +0.j]],\n\n       [[ 6.5+0.j, -1. +0.j],\n        [-0.5+0.j,  0. +0.j]]], dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.ifft2(x, s=[2, 3])\nArray([[[ 1.67+0.j  , -0.08+1.01j, -0.08-1.01j],\n        [-0.33+0.j  , -0.08-0.14j, -0.08+0.14j]],\n\n       [[ 4.33+0.j  ,  0.58+2.17j,  0.58-2.17j],\n        [-0.33+0.j  , -0.08-0.14j, -0.08+0.14j]]], dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.ifft2(x, s=[2, 3], axes=(0, 1))\nArray([[[ 2.33+0.j  ,  3.67+0.j  ],\n        [ 0.33+1.15j,  0.67+1.73j],\n        [ 0.33-1.15j,  0.67-1.73j]],\n\n       [[-1.33+0.j  , -1.33+0.j  ],\n        [-0.33-0.58j, -0.33-0.58j],\n        [-0.33+0.58j, -0.33+0.58j]]], dtype=complex64)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.688Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":40,"estimatedTokens":264}}33{"id":"doc-writing_high_performance_gpu_kernels_with_cute_d-017645fa","source":"documentation","title":"Writing High-Performance GPU Kernels with CuTe DSL and JAX — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/cute_dsl_jax.html","text":"Example:\n```text\noffset = coord[0] * stride[0] + coord[1] * stride[1] + ...\n```\n\nExample:\n```text\ncute.make_layout((...), stride=(...))\n```\n\nExample:\n```text\nrow_major = cute.make_layout((M, N), stride=(N, cutlass.Int32(1)))\ncol_major = cute.make_layout((M, N), stride=(cutlass.Int32(1), M))\n```\n\nExample:\n```text\n!nvidia-smi\n```\n\nExample:\n```text\nimport subprocess\n\n\ndef get_compute_capability():\n  \"\"\"Query the compute capability of the first visible GPU.\"\"\"\n  out = subprocess.check_output(\n      [\"nvidia-smi\", \"--query-gpu=compute_cap\", \"--format=csv,noheader\"],\n      text=True,\n  )\n  major, minor = out.strip().split(\"\\n\")[0].split(\".\")\n  return int(major), int(minor)\n\n\nSM_MAJOR, SM_MINOR = get_compute_capability()\nprint(f\"Detected compute capability: SM {SM_MAJOR}.{SM_MINOR}\")\n\nif SM_MAJOR < 8:\n  print(\"WARNING: CuTe DSL requires SM 8.0+ (Ampere or newer).\")\n  print(\"Some examples may not run on this GPU.\")\nelse:\n  print(\"GPU is compatible with CuTe DSL.\")\n```\n\nExample:\n```text\n%pip install \"nvidia-cutlass-dsl[cu13]\" --quiet\n```\n\nExample:\n```text\nimport os\n\nos.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"2\"  # suppress TF/XLA info & warnings\nos.environ[\"XLA_FLAGS\"] = \"--xla_gpu_cuda_data_dir=/usr/local/cuda\"\n\nimport cutlass\nfrom importlib.metadata import version as _pkg_version\n\nprint(f\"CUTLASS version: {_pkg_version('nvidia-cutlass-dsl')}\")\n\nimport cutlass.cute as cute\nimport cutlass.jax as cjax\nimport cuda.bindings.driver as cuda\n\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\n\nprint(f\"JAX version:     {jax.__version__}\")\nprint(f\"JAX devices:     {jax.devices()}\")\n```\n\nExample:\n```text\n# Optional, if you execute the equivalent kernel definitions further in the notebook\n\n# from cute_dsl_jax.cute_dsl_jax_kernels import (\n#     launch_vector_add, launch_saxpy, launch_gemm,\n#     launch_relu, launch_fused_bias_relu,\n#     launch_elementwise_add,\n# )\n# print(\"Imported: launch_vector_add, launch_saxpy, launch_gemm, launch_relu, launch_fused_bias_relu, launch_elementwise_add\")\n```\n\nExample:\n```text\ndef split_keys(seed=0):\n  key = jax.random.key(seed)\n  while True:\n    key, subkey = jax.random.split(key)\n    yield subkey\n\nkeys = iter(split_keys())\n```\n\nExample:\n```text\ntidx, _, _ = cute.arch.thread_idx()\nbidx, _, _ = cute.arch.block_idx()\n```\n\nExample:\n```text\n@cute.kernel\ndef vector_add_kernel(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):\n  \"\"\"Per-thread kernel: each thread adds one element.\"\"\"\n  tidx, _, _ = cute.arch.thread_idx()\n  bidx, _, _ = cute.arch.block_idx()\n\n  frgA = cute.make_rmem_tensor(cute.size(a, mode=[0]), a.element_type)\n  frgB = cute.make_rmem_tensor(cute.size(b, mode=[0]), b.element_type)\n  frgC = cute.make_rmem_tensor(cute.size(c, mode=[0]), c.element_type)\n\n  cute.autovec_copy(a[None, tidx, bidx], frgA)\n  cute.autovec_copy(b[None, tidx, bidx], frgB)\n  frgC.store(frgA.load() + frgB.load())\n  cute.autovec_copy(frgC, c[None, tidx, bidx])\n```\n\nExample:\n```text\n@cute.jit\ndef launch_vector_add(\n    stream: cuda.CUstream,\n    a: cute.Tensor,\n    b: cute.Tensor,\n    c: cute.Tensor,\n):\n  vector_add_kernel(a, b, c).launch(\n      grid=[a.shape[-1], 1, 1],\n      block=[a.shape[-2], 1, 1],\n      stream=stream,\n  )\n```\n\nExample:\n```text\nN = a.shape[0]\npadded = ((N + BLOCK - 1) // BLOCK) * BLOCK\na_pad = jnp.pad(a, (0, padded - N))\na_3d = a_pad.reshape(1, BLOCK, padded // BLOCK)\n```\n\nExample:\n```text\ncall = cjax.cutlass_call(\n    launch_fn,                    # The @cute.jit function\n    output_shape_dtype=...,       # Shape/dtype of output(s)\n)\nresult = call(*input_arrays)      # Pass JAX arrays here\n```\n\nExample:\n```text\nc_3d = call(a_3d, b_3d)\n```\n\nExample:\n```text\nreturn c_3d.reshape(-1)[:N]\n```\n\nExample:\n```text\nBLOCK = 256  # threads per block for vector add: 256 is a practical default:\n# large enough to expose parallelism, small enough to scale\n# well across different GPUs, and aligned with the hardware’s\n# 32-thread warp execution model.\n\n\n@jax.jit\ndef jax_vector_add(a, b):\n  \"\"\"JAX-compatible vector add using CUTLASS kernel.\"\"\"\n  N = a.shape[0]\n  padded = ((N + BLOCK - 1) // BLOCK) * BLOCK\n  a_pad = jnp.pad(a, (0, padded - N))\n  b_pad = jnp.pad(b, (0, padded - N))\n  # Reshape to (1, BLOCK, num_blocks) for the CuTe kernel\n  a_3d = a_pad.reshape(1, BLOCK, padded // BLOCK)\n  b_3d = b_pad.reshape(1, BLOCK, padded // BLOCK)\n  call = cjax.cutlass_call(\n      launch_vector_add,\n      output_shape_dtype=jax.ShapeDtypeStruct.like(a_3d),\n      use_static_tensors=True,\n  )\n  c_3d = call(a_3d, b_3d)\n  return c_3d.reshape(-1)[:N]\n```\n\nExample:\n```text\n# Test vector add\nN = 1024\na = jax.random.normal(next(keys), (N,), dtype=jnp.float32)\nb = jax.random.normal(next(keys), (N,), dtype=jnp.float32)\n\nc = jax_vector_add(a, b)\nc_ref = a + b\n\nnp.testing.assert_allclose(np.array(c), np.array(c_ref), rtol=1e-5)\nprint(f\"Vector Add PASSED (N={N})\")\nprint(f\"  Max error: {float(jnp.max(jnp.abs(c - c_ref))):.2e}\")\n```\n\nExample:\n```text\ndef saxpy_kernel(x: cute.Tensor, y: cute.Tensor, out: cute.Tensor, alpha: float):\n```\n\nExample:\n```text\nfrgO.store(alpha * frgX.load() + frgY.load())\n```\n\nExample:\n```text\n@cute.kernel\ndef saxpy_kernel(\n    x: cute.Tensor, y: cute.Tensor, out: cute.Tensor, alpha: float\n):\n  \"\"\"SAXPY: out[i] = alpha * x[i] + y[i].\"\"\"\n  tidx, _, _ = cute.arch.thread_idx()\n  bidx, _, _ = cute.arch.block_idx()\n\n  frgX = cute.make_rmem_tensor(cute.size(x, mode=[0]), x.element_type)\n  frgY = cute.make_rmem_tensor(cute.size(y, mode=[0]), y.element_type)\n  frgO = cute.make_rmem_tensor(cute.size(out, mode=[0]), out.element_type)\n\n  cute.autovec_copy(x[None, tidx, bidx], frgX)\n  cute.autovec_copy(y[None, tidx, bidx], frgY)\n  frgO.store(alpha * frgX.load() + frgY.load())\n  cute.autovec_copy(frgO, out[None, tidx, bidx])\n```\n\nExample:\n```text\n@cute.jit\ndef launch_saxpy(\n    stream: cuda.CUstream,\n    x: cute.Tensor,\n    y: cute.Tensor,\n    out: cute.Tensor,\n    *,\n    alpha: float,\n):\n  saxpy_kernel(x, y, out, alpha).launch(\n      grid=[x.shape[-1], 1, 1],\n      block=[x.shape[-2], 1, 1],\n      stream=stream,\n  )\n```\n\nExample:\n```text\ncall = cjax.cutlass_call(\n    launch_saxpy,\n    ...,\n    alpha=alpha,    # scalar kwarg → passed to the kernel\n)\nout_3d = call(x_3d, y_3d)  # tensor args → managed by XLA\n```\n\nExample:\n```text\nBLOCK = 256\n\n\n@jax.jit(static_argnums=(2,))\ndef jax_saxpy(x, y, alpha=2.0):\n  \"\"\"JAX-compatible SAXPY using CUTLASS kernel.\"\"\"\n  N = x.shape[0]\n  padded = ((N + BLOCK - 1) // BLOCK) * BLOCK\n  x_pad = jnp.pad(x, (0, padded - N))\n  y_pad = jnp.pad(y, (0, padded - N))\n  x_3d = x_pad.reshape(1, BLOCK, padded // BLOCK)\n  y_3d = y_pad.reshape(1, BLOCK, padded // BLOCK)\n  call = cjax.cutlass_call(\n      launch_saxpy,\n      output_shape_dtype=jax.ShapeDtypeStruct.like(x_3d),\n      use_static_tensors=True,\n      alpha=alpha,\n  )\n  out_3d = call(x_3d, y_3d)\n  return out_3d.reshape(-1)[:N]\n```\n\nExample:\n```text\n# Test SAXPY\nN = 2048\nALPHA = 2.5\nx = jax.random.normal(next(keys), (N,), dtype=jnp.float32)\ny = jax.random.normal(next(keys), (N,), dtype=jnp.float32)\n\nresult = jax_saxpy(x, y, alpha=ALPHA)\nref = ALPHA * x + y\n\nnp.testing.assert_allclose(np.array(result), np.array(ref),\n                           rtol=1e-5, atol=1e-5)\nprint(f\"SAXPY PASSED (N={N}, alpha={ALPHA})\")\nprint(f\"  Max error: {float(jnp.max(jnp.abs(result - ref))):.2e}\")\n```\n\nExample:\n```text\ntidx, _, _ = cute.arch.thread_idx()\nbidx, _, _ = cute.arch.block_idx()\nbdx, _, _ = cute.arch.block_dim()\n```\n\nExample:\n```text\nidx = bidx * bdx + tidx\n```\n\nExample:\n```text\n@cute.kernel\ndef relu_kernel(x: cute.Tensor, out: cute.Tensor, N: int):\n  \"\"\"Per-thread kernel: each thread computes ReLU of one element.\"\"\"\n  tidx, _, _ = cute.arch.thread_idx()\n  bidx, _, _ = cute.arch.block_idx()\n  bdx, _, _ = cute.arch.block_dim()\n\n  idx = bidx * bdx + tidx\n  if idx < N:\n    val = x[idx]\n    out[idx] = cutlass.max(val, cutlass.Float32(0.0))\n```\n\nExample:\n```text\n@cute.jit\ndef launch_relu(\n    stream: cuda.CUstream,\n    x: cute.Tensor,\n    out: cute.Tensor,\n    *,\n    N: int,\n):\n  BLOCK_SIZE = 256\n  grid_size = (N + BLOCK_SIZE - 1) // BLOCK_SIZE\n  relu_kernel(x, out, N).launch(\n      grid=[grid_size, 1, 1],\n      block=[BLOCK_SIZE, 1, 1],\n      stream=stream,\n  )\n```\n\nExample:\n```text\nx_flat = x.reshape(-1)         # flatten to 1-D\ncall = cjax.cutlass_call(\n    launch_relu,\n    output_shape_dtype=jax.ShapeDtypeStruct.like(x_flat),\n    N=N,                        # scalar kwarg → bounds check inside kernel\n)\nout_flat = call(x_flat)\nreturn out_flat.reshape(x.shape)  # restore original shape\n```\n\nExample:\n```text\n@jax.jit\ndef jax_relu(x):\n  \"\"\"JAX-compatible ReLU using CUTLASS kernel.\"\"\"\n  N = x.size\n  x_flat = x.reshape(-1)\n  call = cjax.cutlass_call(\n      launch_relu,\n      output_shape_dtype=jax.ShapeDtypeStruct.like(x_flat),\n      N=N,\n  )\n  out_flat = call(x_flat)\n  return out_flat.reshape(x.shape)\n```\n\nExample:\n```text\n# Test ReLU\nN = 2048\nx = jax.random.normal(next(keys), (N,), dtype=jnp.float32)\n\nresult = jax_relu(x)\nref = jax.nn.relu(x)\n\nnp.testing.assert_allclose(np.array(result), np.array(ref), rtol=1e-5)\nprint(f\"ReLU PASSED (N={N})\")\nprint(f\"  Max error: {float(jnp.max(jnp.abs(result - ref))):.2e}\")\nprint(f\"  Sample: x[:6]   = {x[:6]}\")\nprint(f\"          out[:6] = {result[:6]}\")\n```\n\nExample:\n```text\n@cute.kernel\ndef fused_bias_relu_kernel(\n    x: cute.Tensor,\n    bias: cute.Tensor,\n    out: cute.Tensor,\n    N: int,\n    width: int,\n):\n  \"\"\"Per-thread: out[i] = max(0, x[i] + bias[i % width]).\"\"\"\n  tidx, _, _ = cute.arch.thread_idx()\n  bidx, _, _ = cute.arch.block_idx()\n  bdx, _, _ = cute.arch.block_dim()\n\n  idx = bidx * bdx + tidx\n  if idx < N:\n    col = idx % width\n    val = x[idx] + bias[col]\n    out[idx] = cutlass.max(val, cutlass.Float32(0.0))\n```\n\nExample:\n```text\n@cute.jit\ndef launch_fused_bias_relu(\n    stream: cuda.CUstream,\n    x: cute.Tensor,\n    bias: cute.Tensor,\n    out: cute.Tensor,\n    *,\n    N: int,\n    width: int,\n):\n  BLOCK_SIZE = 256\n  grid_size = (N + BLOCK_SIZE - 1) // BLOCK_SIZE\n  fused_bias_relu_kernel(x, bias, out, N, width).launch(\n      grid=[grid_size, 1, 1],\n      block=[BLOCK_SIZE, 1, 1],\n      stream=stream,\n  )\n```\n\nExample:\n```text\ncall = cjax.cutlass_call(\n    launch_fused_bias_relu,\n    output_shape_dtype=jax.ShapeDtypeStruct.like(x_flat),\n    N=N, width=width,\n)\nout_flat = call(x_flat, bias)   # two input tensors: x and bias\n```\n\nExample:\n```text\n@jax.jit(static_argnums=(2,))\ndef jax_fused_bias_relu(x, bias, width):\n  \"\"\"JAX-compatible fused Bias+ReLU using CUTLASS kernel.\n\n  Args:\n      x: Input matrix of shape (batch, width), flattened to 1-D for the kernel.\n      bias: Bias vector of shape (width,).\n      width: Number of columns (static, passed as constexpr to the kernel).\n  \"\"\"\n  N = x.size\n  x_flat = x.reshape(-1)\n  call = cjax.cutlass_call(\n      launch_fused_bias_relu,\n      output_shape_dtype=jax.ShapeDtypeStruct.like(x_flat),\n      N=N,\n      width=width,\n  )\n  out_flat = call(x_flat, bias)\n  return out_flat.reshape(x.shape)\n```\n\nExample:\n```text\n# Test Fused Bias+ReLU\nBATCH, WIDTH = 64, 512\nx = jax.random.normal(next(keys), (BATCH, WIDTH), dtype=jnp.float32)\nbias = jax.random.normal(next(keys), (WIDTH,), dtype=jnp.float32)\n\nresult = jax_fused_bias_relu(x, bias, WIDTH)\nref = jnp.maximum(0, x + bias[None, :])\n\nnp.testing.assert_allclose(np.array(result), np.array(ref), rtol=1e-5)\nprint(f\"Fused Bias+ReLU PASSED (batch={BATCH}, width={WIDTH})\")\nprint(f\"  Max error: {float(jnp.max(jnp.abs(result - ref))):.2e}\")\n```\n\nExample:\n```text\ntidx, _, _ = cute.arch.thread_idx()\n    bm, bn, _ = cute.arch.block_idx()\n    bdx, _, _ = cute.arch.block_dim()\n```\n\nExample:\n```text\nfor i in cutlass.range(tidx, BLOCK_M * BLOCK_N, bdx):\n```\n\nExample:\n```text\nrow = i // BLOCK_N          # tile-local row\n        col = i % BLOCK_N           # tile-local column\n        m_idx = bm * BLOCK_M + row  # global row in D\n        n_idx = bn * BLOCK_N + col  # global column in D\n```\n\nExample:\n```text\nif m_idx < M and n_idx < N:\n```\n\nExample:\n```text\nacc = cutlass.Float32(0.0)\n            for k in cutlass.range(K):\n                acc += A[m_idx * K + k] * B[k * N + n_idx]\n            D[m_idx * N + n_idx] = acc\n```\n\nExample:\n```text\n@cute.kernel\ndef gemm_kernel(\n    A: cute.Tensor,\n    B: cute.Tensor,\n    D: cute.Tensor,\n    M: int,\n    N: int,\n    K: int,\n    BLOCK_M: int,\n    BLOCK_N: int,\n):\n  \"\"\"Tiled GEMM: each thread accumulates output elements.\"\"\"\n  tidx, _, _ = cute.arch.thread_idx()\n  bm, bn, _ = cute.arch.block_idx()\n  bdx, _, _ = cute.arch.block_dim()\n\n  for i in cutlass.range(tidx, BLOCK_M * BLOCK_N, bdx):\n    row = i // BLOCK_N\n    col = i % BLOCK_N\n    m_idx = bm * BLOCK_M + row\n    n_idx = bn * BLOCK_N + col\n    if m_idx < M and n_idx < N:\n      acc = cutlass.Float32(0.0)\n      for k in cutlass.range(K):\n        acc += A[m_idx * K + k] * B[k * N + n_idx]\n      D[m_idx * N + n_idx] = acc\n```\n\nExample:\n```text\n@cute.jit\ndef launch_gemm(\n    stream: cuda.CUstream,\n    A: cute.Tensor,\n    B: cute.Tensor,\n    D: cute.Tensor,\n    *,\n    M: int,\n    N: int,\n    K: int,\n):\n  BLOCK_M, BLOCK_N = 64, 64\n  grid_m = (M + BLOCK_M - 1) // BLOCK_M\n  grid_n = (N + BLOCK_N - 1) // BLOCK_N\n  gemm_kernel(A, B, D, M, N, K, BLOCK_M, BLOCK_N).launch(\n      grid=[grid_m, grid_n, 1],\n      block=[256, 1, 1],\n      stream=stream,\n  )\n```\n\nExample:\n```text\n@jax.jit\ndef jax_cutlass_gemm(a, b):\n  \"\"\"JAX wrapper for the CUTLASS GEMM kernel.\"\"\"\n  M, K = a.shape\n  _, N = b.shape\n  a_flat = a.reshape(-1)\n  b_flat = b.reshape(-1)\n  call = cjax.cutlass_call(\n      launch_gemm,\n      output_shape_dtype=jax.ShapeDtypeStruct((M * N,), a.dtype),\n      M=M,\n      N=N,\n      K=K,\n  )\n  d_flat = call(a_flat, b_flat)\n  return d_flat.reshape(M, N)\n```\n\nExample:\n```text\n# Test GEMM\nM, N, K = 256, 256, 128\nA = jax.random.normal(next(keys), (M, K), dtype=jnp.float32)\nB = jax.random.normal(next(keys), (K, N), dtype=jnp.float32)\n\nD = jax_cutlass_gemm(A, B)\nD_ref = jnp.matmul(A, B)\n\nnp.testing.assert_allclose(np.array(D), np.array(D_ref), rtol=1e-2, atol=2e-2)\nprint(f\"GEMM PASSED (M={M}, N={N}, K={K})\")\nprint(f\"  Max error: {float(jnp.max(jnp.abs(D - D_ref))):.2e}\")\n```\n\nExample:\n```text\nimport time\n\nM, N, K = 512, 512, 512\nA = jax.random.normal(next(keys), (M, K), dtype=jnp.float32)\nB = jax.random.normal(next(keys), (K, N), dtype=jnp.float32)\n\n# Warmup\n_ = jax_cutlass_gemm(A, B).block_until_ready()\n_ = jnp.matmul(A, B).block_until_ready()\n\nNUM_RUNS = 20\n\n# Time CUTLASS GEMM\nstart = time.perf_counter()\nfor _ in range(NUM_RUNS):\n  _ = jax_cutlass_gemm(A, B).block_until_ready()\ncutlass_time = (time.perf_counter() - start) / NUM_RUNS\n\n# Time JAX matmul\nstart = time.perf_counter()\nfor _ in range(NUM_RUNS):\n  _ = jnp.matmul(A, B).block_until_ready()\njax_time = (time.perf_counter() - start) / NUM_RUNS\n\nprint(f\"Matrix size: {M}x{N}x{K}\")\nprint(f\"CUTLASS GEMM:  {cutlass_time*1000:.3f} ms\")\nprint(f\"JAX jnp.matmul: {jax_time*1000:.3f} ms\")\nprint(f\"Ratio (CUTLASS / JAX): {cutlass_time / jax_time:.2f}x\")\nprint()\nprint(\"Note: Our simple tiled kernel is not expected to beat cuBLAS.\")\nprint(\"CuTe DSL's value is in specialized kernels cuBLAS doesn't provide.\")\n```\n\nExample:\n```text\nmesh = jax.make_mesh((num_devices,), (\"x\",))\njax.set_mesh(mesh)\n```\n\nExample:\n```text\nsharding = P(None, None, \"x\")\n```\n\nExample:\n```text\na = jax.random.normal(\n    jax.random.key(10),\n    shape,\n    dtype=jnp.float32,\n    out_sharding=sharding,\n)\nb = jax.random.normal(\n    jax.random.key(11),\n    shape,\n    dtype=jnp.float32,\n    out_sharding=sharding,\n)\n```\n\nExample:\n```text\nfrom jax.sharding import NamedSharding\n\nnamed_sharding = NamedSharding(mesh, sharding)\n\na = jax.random.normal(jax.random.key(10), shape, dtype=jnp.float32)\nb = jax.random.normal(jax.random.key(11), shape, dtype=jnp.float32)\n\na = jax.device_put(a, named_sharding)\nb = jax.device_put(b, named_sharding)\n```\n\nExample:\n```text\n@jax.shard_map(out_specs=sharding)\ndef sharded_vector_add(a_shard, b_shard):\n    call = cjax.cutlass_call(\n        launch_vector_add,\n        output_shape_dtype=jax.typeof(a_shard),\n        use_static_tensors=True,\n    )\n    return call(a_shard, b_shard)\n```\n\nExample:\n```text\nfrom jax.sharding import PartitionSpec as P\n\nnum_devices = len(jax.devices())\nprint(f\"Number of devices: {num_devices}\")\n\nBLOCK = 256\n\nmesh = jax.make_mesh((num_devices,), (\"x\",))\n\n# Use `jax.set_mesh` as a context manager so the mesh is scoped to this\n# sharding demo and does not leak into later cells.\nwith jax.set_mesh(mesh):\n  # Kernel expects 3-D tensors: (elems_per_thread, threads, blocks)\n  # Shard along the blocks axis (last dim)\n  sharding = P(None, None, \"x\")\n\n  @jax.shard_map(out_specs=sharding)\n  def sharded_vector_add(a_shard, b_shard):\n    call = cjax.cutlass_call(\n        launch_vector_add,\n        output_shape_dtype=jax.typeof(a_shard),\n        use_static_tensors=True,\n    )\n    return call(a_shard, b_shard)\n\n  # Create 3-D tensors: (1, 256, total_blocks) with total_blocks divisible by device count\n  blocks_per_device = 16\n  total_blocks = blocks_per_device * num_devices\n  shape = (1, BLOCK, total_blocks)\n\n  a_m = jax.random.normal(\n      jax.random.key(10),\n      shape,\n      dtype=jnp.float32,\n      out_sharding=sharding,\n  )\n  b_m = jax.random.normal(\n      jax.random.key(11),\n      shape,\n      dtype=jnp.float32,\n      out_sharding=sharding,\n  )\n\n  print(\"a_m sharding:\", a_m.sharding)\n  print(\"b_m sharding:\", b_m.sharding)\n\n  c_m = sharded_vector_add(a_m, b_m)\n\n  np.testing.assert_allclose(jnp.array(c_m), jnp.array(a_m + b_m), rtol=1e-5)\n  n_total = int(np.prod(shape))\n  print(f\"Sharded Vector Add PASSED across {num_devices} devices (N={n_total})\")\n```\n\nExample:\n```text\nfrom jax import export\nfrom cutlass.jax import get_export_disabled_safety_checks\n\n# 1. Export the JIT-compiled function\nexported = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks())\n\n# 2. Specialize to a signature (concrete or symbolic shapes) and serialize\ntraced = exported(shape_dtype_spec, shape_dtype_spec)\nblob = traced.serialize()\n\n# 3. Later: deserialize and call with real data\nrehydrated = export.deserialize(blob)\nresult = rehydrated.call(a, b)\n```\n\nExample:\n```text\nfrom cutlass.jax import get_export_disabled_safety_checks\nfrom jax import export\n\n\n# Element-wise Add (2-D, flat indexing)\n@cute.kernel\ndef elementwise_add_kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):\n  \"\"\"Per-thread kernel: 2-D element-wise add using flat indexing.\"\"\"\n  tidx, _, _ = cute.arch.thread_idx()\n  bidx, _, _ = cute.arch.block_idx()\n  bdim, _, _ = cute.arch.block_dim()\n\n  thread_idx = bidx * bdim + tidx\n\n  m, n = gA.shape\n\n  if thread_idx < m * n:\n\n    ni = thread_idx % n\n    mi = thread_idx // n\n\n    a_val = gA[mi, ni]\n    b_val = gB[mi, ni]\n    gC[mi, ni] = a_val + b_val\n\n\n@cute.jit\ndef launch_elementwise_add(\n    stream: cuda.CUstream,\n    mA: cute.Tensor,\n    mB: cute.Tensor,\n    mC: cute.Tensor,\n):\n  num_threads_per_block = 256\n  m, n = mA.shape\n  elementwise_add_kernel(mA, mB, mC).launch(\n      grid=((m * n + num_threads_per_block - 1) // num_threads_per_block, 1, 1),\n      block=(num_threads_per_block, 1, 1),\n      stream=stream,\n  )\n\n\n# Define a function that uses a CUTLASS kernel + JAX ops.\n# We use launch_elementwise_add which accepts 2-D tensors directly\n# with flat indexing — compatible with jax.export's tracing.\n@jax.jit\ndef f(a, b):\n  call = cjax.cutlass_call(launch_elementwise_add, output_shape_dtype=a)\n  return jax.nn.sigmoid(call(a, b))\n\n\n# Reference implementation (pure JAX)\n@jax.jit\ndef ref_f(a, b):\n  return jax.nn.sigmoid(a + b)\n\n\n# --- Export with concrete shapes ---\nM, N = 512, 256\nexport_shape_dtype = jax.ShapeDtypeStruct((M, N), jnp.float32)\n\nprint(\n    f\"Exporting with input signature: ({export_shape_dtype},\"\n    f\" {export_shape_dtype})\"\n)\n\n# Export the function — get_export_disabled_safety_checks() tells JAX\n# that CUTLASS custom call targets are safe to include\nexported = jax.export.export(\n    f, disabled_checks=get_export_disabled_safety_checks()\n)\ntraced = exported(export_shape_dtype, export_shape_dtype)\n\n# Serialize to a byte blob\nblob = traced.serialize()\nprint(f\"Serialized computation: {len(blob):,} bytes\")\n\n# Deserialize and run — this works independently of the original function\nrehydrated = export.deserialize(blob)\n\na = jax.random.normal(next(keys), (M, N), dtype=jnp.float32)\nb = jax.random.normal(next(keys), (M, N), dtype=jnp.float32)\n\nc = rehydrated.call(a, b)\nc_ref = ref_f(a, b)\n\nnp.testing.assert_allclose(np.array(c), np.array(c_ref), rtol=1e-5)\nprint(f\"Export + Deserialize PASSED (M={M}, N={N})\")\nprint(f\"  Max error: {float(jnp.max(jnp.abs(c - c_ref))):.2e}\")\n```\n\nExample:\n```text\n# --- Export with symbolic shapes ---\na_sym, b_sym = export.symbolic_shape(\"a, b\")\nsymbolic_shape_dtype = jax.ShapeDtypeStruct((a_sym, b_sym), jnp.float32)\n\nprint(\n    f\"Exporting with symbolic signature: ({symbolic_shape_dtype},\"\n    f\" {symbolic_shape_dtype})\"\n)\n\nexported_sym = jax.export.export(\n    f, disabled_checks=get_export_disabled_safety_checks()\n)\ntraced_sym = exported_sym(symbolic_shape_dtype, symbolic_shape_dtype)\nblob_sym = traced_sym.serialize()\nprint(f\"Serialized computation: {len(blob_sym):,} bytes\")\n\nrehydrated_sym = export.deserialize(blob_sym)\n\n# Call with different shapes — no recompilation needed.\n# The same serialized blob works for any (M, N) where M*N is a\n# multiple of the kernel's block size (256).\nfor shape in [(512, 256), (1024, 512), (2048, 1024)]:\n  a = jax.random.normal(next(keys), shape, dtype=jnp.float32)\n  b = jax.random.normal(next(keys), shape, dtype=jnp.float32)\n  c = rehydrated_sym.call(a, b)\n  c_ref = ref_f(a, b)\n  np.testing.assert_allclose(np.array(c), np.array(c_ref), rtol=1e-5)\n  print(f\"  Symbolic export PASSED for shape {shape}\")\n\nprint(\"All symbolic shape tests passed.\")\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.691Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":56,"totalLines":895,"estimatedTokens":5428}}34{"id":"doc-defining_new_jax_types_with_hijax_jax_documentat-abbb805f","source":"documentation","title":"Defining new JAX types with hijax — JAX documentation","url":"https://docs.jax.dev/en/latest/hijax_types.html","text":"Example:\n```text\nimport os\nos.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=8'\n# (8 CPU devices, for the sharding sections at the end)\n\nfrom dataclasses import dataclass\n\nimport jax\nimport jax.numpy as jnp\n\n@dataclass(frozen=True)\nclass QArray:\n  qvalue: jax.Array   # int8[*leading, n]\n  scale: jax.Array    # f32[*leading]\n```\n\nExample:\n```text\nfrom jax.experimental.hijax import HiType, ShapedArray, register_hitype\nfrom jax.sharding import NamedSharding\n\n@dataclass(frozen=True)\nclass QArrayTy(HiType):\n  shape: tuple[int, ...]\n  sharding: NamedSharding  # qvalue's sharding; scale's is derived from it\n\n  # lowering: which array types make up this type, and how values convert\n  def lo_ty(self):\n    scale_sharding = self.sharding.update(spec=jax.P(*self.sharding.spec[:-1]))\n    return [ShapedArray(self.shape, jnp.dtype('int8'),\n                        sharding=self.sharding),\n            ShapedArray(self.shape[:-1], jnp.dtype('float32'),\n                        sharding=scale_sharding)]\n  def lower_val(self, q):\n    return [q.qvalue, q.scale]\n  def raise_val(self, qvalue, scale):\n    return QArray(qvalue, scale)\n\n  # autodiff: tangents of quantized arrays are plain float arrays (see below)\n  def to_tangent_aval(self):\n    return ShapedArray(self.shape, jnp.dtype('float32'),\n                       sharding=self.sharding)\n\n  # printing, e.g. in jaxprs\n  def str_short(self, short_dtypes=False, mesh_axis_types=False):\n    dims = [str(d) if p is None else f'{d}@{p}'\n            for d, p in zip(self.shape, self.sharding.spec)]\n    return f'q8[{\",\".join(dims)}]'\n  __repr__ = str_short\n\nregister_hitype(QArray, lambda q: QArrayTy(q.qvalue.shape,\n                                           jax.typeof(q.qvalue).sharding))\n```\n\nExample:\n```text\nfrom jax.experimental.hijax import VJPHiPrimitive\n\nclass Quantize(VJPHiPrimitive):\n  def __init__(self, x_aval):\n    if x_aval.dtype != jnp.dtype('float32'): raise TypeError(x_aval.dtype)\n    self.in_avals = (x_aval,)\n    self.out_aval = QArrayTy(x_aval.shape, x_aval.sharding)\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x):\n    scale = jnp.max(jnp.abs(x), axis=-1) / 127.\n    qvalue = jnp.round(x / scale[..., None]).astype(jnp.int8)\n    return QArray(qvalue, scale)\n\n  # straight-through estimator: differentiate as if it's the identity\n  def vjp_fwd(self, nzs_in, x):\n    return self(x), None\n\n  def vjp_bwd_retval(self, _res, g):\n    return (g,)\n\nclass Dequantize(VJPHiPrimitive):\n  def __init__(self, q_aval):\n    self.in_avals = (q_aval,)\n    self.out_aval = ShapedArray(q_aval.shape, jnp.dtype('float32'),\n                                sharding=q_aval.sharding)\n    self.params = {}\n    super().__init__()\n\n  def expand(self, qx):\n    return qx.qvalue.astype('float32') * qx.scale[..., None]\n\n  def vjp_fwd(self, nzs_in, qx):\n    return self(qx), None\n\n  def vjp_bwd_retval(self, _res, g):\n    return (g,)\n\ndef quantize(x):\n  return Quantize(jax.typeof(x))(x)\n\ndef dequantize(qx):\n  return Dequantize(jax.typeof(qx))(qx)\n```\n\nExample:\n```text\nx = jnp.array([[1., 2., 3.],\n               [4., -5., 6.]])\n\nqx = quantize(x)\nprint(qx)\nprint(jax.typeof(qx))\nprint(dequantize(qx))\n```\n\nExample:\n```text\nQArray(qvalue=Array([[  42,   85,  127],\n       [  85, -106,  127]], dtype=int8), scale=Array([0.02362205, 0.04724409], dtype=float32))\nq8[2,3]\n[[ 0.992126  2.007874  3.      ]\n [ 4.015748 -5.007874  6.      ]]\n```\n\nExample:\n```text\ntry:\n  jax.jit(lambda qx: qx.qvalue)(qx)\nexcept AttributeError as e:\n  print('AttributeError:', e)\n```\n\nExample:\n```text\nAttributeError: DynamicJaxprTracer has no attribute qvalue\n```\n\nExample:\n```text\ndef bad_quantize(x):\n  scale = jnp.max(jnp.abs(x), axis=-1) / 127.\n  return QArray(jnp.round(x / scale[..., None]).astype('int8'), scale)\n\ntry:\n  jax.jit(bad_quantize)(x)\nexcept TypeError as e:\n  print('TypeError:', e)\n```\n\nExample:\n```text\nTypeError: No constant handler for type: <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>\n```\n\nExample:\n```text\nclass MatmulQ(VJPHiPrimitive):\n  def __init__(self, x_aval, q_aval):\n    if not (isinstance(q_aval, QArrayTy) and len(x_aval.shape) == 2 and\n            len(q_aval.shape) == 2 and x_aval.shape[1] == q_aval.shape[0]):\n      raise TypeError(f'bad matmul_q operand types: {x_aval} @ {q_aval}')\n    self.in_avals = (x_aval, q_aval)\n    x_spec, q_spec = x_aval.sharding.spec, q_aval.sharding.spec\n    if x_spec[1] is not None or q_spec[0] is not None:\n      raise TypeError('matmul_q requires unsharded contraction axes, got '\n                      f'{x_aval} @ {q_aval}')\n    out_sharding = x_aval.sharding.update(spec=jax.P(x_spec[0], q_spec[1]))\n    self.out_aval = ShapedArray((x_aval.shape[0], q_aval.shape[1]),\n                                jnp.dtype('float32'), sharding=out_sharding)\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x, qw):\n    # fold the per-row scales into the dense operand, then apply one matmul\n    # directly against the int8 payload\n    return (x * qw.scale) @ qw.qvalue.astype(jnp.float32)\n\n  def vjp_fwd(self, nzs_in, x, qw):\n    return self(x, qw), (x, qw)\n\n  def vjp_bwd_retval(self, res, g):\n    x, qw = res\n    w = dequantize(qw)   # rules are traced code: use primitives here\n    # cotangents live where the primals live, so use the primal operands'\n    # shardings to disambiguate the (possibly sharded) contractions\n    return (jnp.matmul(g, w.T, out_sharding=jax.typeof(x).sharding),\n            jnp.matmul(x.T, g, out_sharding=jax.typeof(w).sharding))\n\ndef matmul_q(x, qw):\n  return MatmulQ(jax.typeof(x), jax.typeof(qw))(x, qw)\n```\n\nExample:\n```text\nw = jnp.arange(12., dtype='float32').reshape(3, 4) / 12.\nqw = quantize(w)\n\nprint(matmul_q(x, qw))\nprint(x @ dequantize(qw))  # reference\n```\n\nExample:\n```text\n[[2.6627297 3.1706038 3.6587927 4.166667 ]\n [2.3077428 2.7447505 3.1463253 3.5833333]]\n[[2.6627297 3.1706038 3.6587927 4.166667 ]\n [2.3077426 2.7447505 3.1463253 3.5833333]]\n```\n\nExample:\n```text\njax.jit(lambda x: dequantize(quantize(x))).trace(x).jaxpr\n```\n\nExample:\n```text\n{ lambda ; a:f32[2,3]. let\n    b:q8[2,3] = call_hi_primitive[_prim=Quantize[{}]] a\n    c:f32[2,3] = call_hi_primitive[_prim=Dequantize[{}]] b\n  in (c,) }\n```\n\nExample:\n```text\njax.jit(matmul_q).trace(x, qw).jaxpr\n#\n# `jit` works, with quantized arrays as arguments, results, and\n# intermediates:\n```\n\nExample:\n```text\n{ lambda ; a:f32[2,3] b:q8[3,4]. let\n    c:f32[2,4] = call_hi_primitive[_prim=MatmulQ[{}]] a b\n  in (c,) }\n```\n\nExample:\n```text\nprint(jax.jit(lambda x: dequantize(quantize(x)))(x))   # QArray internal\n\nqx2 = jax.jit(quantize)(x)                             # QArray result\nprint(jax.typeof(qx2))\n\nprint(jax.jit(dequantize)(qx2))                        # QArray argument\n```\n\nExample:\n```text\n[[ 0.992126  2.007874  3.      ]\n [ 4.015748 -5.007874  6.      ]]\nq8[2,3]\n[[ 0.992126  2.007874  3.      ]\n [ 4.015748 -5.007874  6.      ]]\n```\n\nExample:\n```text\ndef to_tangent_aval(self):\n    return ShapedArray(self.shape, jnp.dtype('float32'))\n```\n\nExample:\n```text\ndef f(x):\n  return jnp.sum(dequantize(quantize(x)))\n```\n\nExample:\n```text\nprint(jax.grad(f)(x))\n```\n\nExample:\n```text\n[[1. 1. 1.]\n [1. 1. 1.]]\n```\n\nExample:\n```text\ndef g(qx):\n  return jnp.sum(dequantize(qx) ** 2)\n\nprint(jax.grad(g)(qx))\nprint(jax.typeof(jax.grad(g)(qx)))\n```\n\nExample:\n```text\n[[  1.984252   4.015748   6.      ]\n [  8.031496 -10.015748  12.      ]]\nfloat32[2,3]\n```\n\nExample:\n```text\ndef loss(x, qw):\n  return jnp.sum(matmul_q(x, qw) ** 2)\n\ngrad_x, grad_qw = jax.grad(loss, argnums=(0, 1))(x, qw)\nprint(jax.typeof(grad_x), jax.typeof(grad_qw))\n```\n\nExample:\n```text\nfloat32[2,3] float32[3,4]\n```\n\nExample:\n```text\nfrom jax.experimental.hijax import MappingSpec\n\n@dataclass(frozen=True)\nclass QArraySpec(MappingSpec):\n  pass  # QArrays are only mapped along their leading axis\n```\n\nExample:\n```text\ndef qarray_dec_rank(self, size, spec):\n  assert isinstance(spec, QArraySpec) and self.shape[0] == size\n  return QArrayTy(self.shape[1:],\n                  self.sharding.update(spec=jax.P(*self.sharding.spec[1:])))\n\ndef qarray_inc_rank(self, size, spec):\n  assert isinstance(spec, QArraySpec)\n  return QArrayTy((size, *self.shape),\n                  self.sharding.update(spec=jax.P(None, *self.sharding.spec)))\n\nQArrayTy.dec_rank = qarray_dec_rank\nQArrayTy.inc_rank = qarray_inc_rank\n```\n\nExample:\n```text\ndef quantize_batch(self, axis_data, args, in_dims):\n  x, = args\n  d, = in_dims\n  if d is None:\n    return quantize(x), None\n  x = jnp.moveaxis(x, d, 0)\n  return quantize(x), QArraySpec()\nQuantize.batch = quantize_batch\n\ndef dequantize_batch(self, axis_data, args, in_dims):\n  qx, = args\n  d, = in_dims\n  if d is None:\n    return dequantize(qx), None\n  assert isinstance(d, QArraySpec)\n  return dequantize(qx), 0\nDequantize.batch = dequantize_batch\n```\n\nExample:\n```text\nxs = jnp.arange(24., dtype='float32').reshape(4, 2, 3)\n\nqxs = jax.vmap(quantize, out_axes=QArraySpec())(xs)\nprint(jax.typeof(qxs))\nprint(qxs.qvalue.shape, qxs.scale.shape)\n```\n\nExample:\n```text\nq8[4,2,3]\n\n(4, 2, 3) (4, 2)\n```\n\nExample:\n```text\nxs_roundtrip = jax.vmap(dequantize, in_axes=QArraySpec(), axis_size=4)(qxs)\nprint(jax.typeof(xs_roundtrip))\n```\n\nExample:\n```text\nfloat32[4,2,3]\n```\n\nExample:\n```text\nprint(jax.typeof(jax.vmap(jax.jit(dequantize), in_axes=QArraySpec(),\n                          axis_size=4)(qxs)))\n```\n\nExample:\n```text\ndef norm_quantized(x):\n  return jnp.sum(dequantize(quantize(x)) ** 2)\n\nprint(jax.vmap(jax.grad(norm_quantized))(xs).shape)\n```\n\nExample:\n```text\n(4, 2, 3)\n```\n\nExample:\n```text\ndef qarray_leading_axis_spec(self):\n  return QArraySpec()\n\nQArrayTy.leading_axis_spec = qarray_leading_axis_spec\n```\n\nExample:\n```text\ndef sum_dequantized(total, qx):\n  return total + jnp.sum(dequantize(qx)), ()\n\ntotal, () = jax.lax.scan(sum_dequantized, 0., qxs, length=4)\nprint(total)\n```\n\nExample:\n```text\n275.9685\n```\n\nExample:\n```text\ndef accum_quantized(qtotal, x):\n  return quantize(dequantize(qtotal) + x), quantize(2 * x)\n\nqzero = quantize(jnp.zeros((2, 3), 'float32'))\nqtotal, qys = jax.lax.scan(accum_quantized, qzero, xs)\nprint(jax.typeof(qtotal))\nprint(jax.typeof(qys))\n```\n\nExample:\n```text\nq8[2,3]\nq8[4,2,3]\n```\n\nExample:\n```text\nmesh = jax.make_mesh((4,), ('i',))\njax.set_mesh(mesh)\n\nrows = jax.device_put(jnp.arange(24., dtype='float32').reshape(8, 3),\n                      jax.P('i'))\n\nqrows = quantize(rows)\nprint(jax.typeof(qrows))\nprint(qrows.qvalue.sharding.spec, qrows.scale.sharding.spec)\nprint(jax.typeof(dequantize(qrows)))\n```\n\nExample:\n```text\nq8[8@i,3]\nP('i', None) P('i',)\nfloat32[8@i,3]\n```\n\nExample:\n```text\njax.jit(lambda x: dequantize(quantize(x))).trace(rows).jaxpr\n```\n\nExample:\n```text\n{ lambda ; a:f32[8@i,3]. let\n    b:q8[8@i,3] = call_hi_primitive[_prim=Quantize[{}]] a\n    c:f32[8@i,3] = call_hi_primitive[_prim=Dequantize[{}]] b\n  in (c,) }\n```\n\nExample:\n```text\nw2 = jnp.arange(12., dtype='float32').reshape(3, 4) / 12.\n\nprint(jax.typeof(matmul_q(rows, quantize(w2))))                # rows sharded\n\nqw2 = quantize(jax.device_put(w2, jax.P(None, 'i')))\nprint(jax.typeof(qw2))                                         # cols sharded\nprint(jax.typeof(matmul_q(jnp.ones((2, 3), 'float32'), qw2)))\n```\n\nExample:\n```text\nfloat32[8@i,4]\nq8[3,4@i]\nfloat32[2,4@i]\n```\n\nExample:\n```text\nxk = jax.device_put(jnp.ones((2, 8), 'float32'), jax.P(None, 'i'))\nqk = quantize(jax.device_put(jnp.ones((8, 3), 'float32'), jax.P('i')))\n\ntry:\n  matmul_q(xk, qk)\nexcept TypeError as e:\n  print('TypeError:', e)\n```\n\nExample:\n```text\nTypeError: matmul_q requires unsharded contraction axes, got float32[2,8@i] @ q8[8@i,3]\n```\n\nExample:\n```text\ndef qloss(x, qw):\n  return jnp.sum(matmul_q(x, qw) ** 2)\n\ngrad_rows, grad_qw2 = jax.grad(qloss, argnums=(0, 1))(rows, quantize(w2))\nprint(jax.typeof(grad_rows), jax.typeof(grad_qw2))\n```\n\nExample:\n```text\nfloat32[8@i,3] float32[3,4]\n```\n\nExample:\n```text\nfrom jax.experimental.hijax import HiPspec\n\n@dataclass(frozen=True)\nclass QArrayP(HiPspec):\n  spec: jax.P  # partitioning of the leading axes; the last axis stays whole\n\n  def to_lo(self):\n    return (self.spec, self.spec)  # qvalue and scale shard together\n```\n\nExample:\n```text\ndef qarray_shard(self, mesh, manual_axes, check_vma, spec):\n  qvalue_ty, _ = self.lo_ty()\n  qspec, _ = spec.to_lo()\n  shard_ty = qvalue_ty.shard(mesh, manual_axes, check_vma, qspec)\n  return QArrayTy(shard_ty.shape, shard_ty.sharding)\n\ndef qarray_unshard(self, mesh, check_vma, spec):\n  qvalue_ty, _ = self.lo_ty()\n  qspec, _ = spec.to_lo()\n  full_ty = qvalue_ty.unshard(mesh, check_vma, qspec)\n  return QArrayTy(full_ty.shape, full_ty.sharding)\n\nQArrayTy.shard = qarray_shard\nQArrayTy.unshard = qarray_unshard\n```\n\nExample:\n```text\n@jax.jit\n@jax.shard_map(in_specs=jax.P('i'), out_specs=QArrayP(jax.P('i')))\ndef quantize_shards(x):\n  assert jax.typeof(x).shape == (2, 3)   # each device sees two rows\n  return quantize(x)\n\nqrows = quantize_shards(rows)\nprint(jax.typeof(qrows))\nprint(qrows.qvalue.sharding.spec, qrows.scale.sharding.spec)\n```\n\nExample:\n```text\nq8[8@i,3]\nP('i', None) P('i',)\n```\n\nExample:\n```text\n@jax.jit\n@jax.shard_map(in_specs=QArrayP(jax.P('i')), out_specs=jax.P('i'))\ndef dequantize_shards(qx):\n  assert jax.typeof(qx).shape == (2, 3)  # a per-device QArray shard\n  return dequantize(qx)\n\nprint(jnp.max(jnp.abs(dequantize_shards(qrows) - rows)))\n```\n\nExample:\n```text\n0.08661461\n```\n\nExample:\n```text\nqrows_global = quantize(rows)\nassert (qrows.qvalue == qrows_global.qvalue).all()\nassert (qrows.scale == qrows_global.scale).all()\n```\n\nExample:\n```text\n@dataclass(frozen=True)\nclass Rank1:\n  col: jax.Array   # f32[m]\n  row: jax.Array   # f32[n]\n\n@dataclass(frozen=True)\nclass Rank1Ty(HiType):\n  shape: tuple[int, int]   # (m, n), the dense shape represented\n  sharding: NamedSharding  # sharding of the dense shape; the factors' derive\n\n  def lo_ty(self):\n    (m, n), spec = self.shape, self.sharding.spec\n    return [ShapedArray((m,), jnp.dtype('float32'),\n                        sharding=self.sharding.update(spec=jax.P(spec[0]))),\n            ShapedArray((n,), jnp.dtype('float32'),\n                        sharding=self.sharding.update(spec=jax.P(spec[1])))]\n  def lower_val(self, r1):\n    return [r1.col, r1.row]\n  def raise_val(self, col, row):\n    return Rank1(col, row)\n\n  # rank-1 matrices aren't closed under addition, so tangents are dense\n  def to_tangent_aval(self):\n    return ShapedArray(self.shape, jnp.dtype('float32'),\n                       sharding=self.sharding)\n\n  def str_short(self, short_dtypes=False, mesh_axis_types=False):\n    dims = [str(d) if p is None else f'{d}@{p}'\n            for d, p in zip(self.shape, self.sharding.spec)]\n    return f'r1[{\",\".join(dims)}]'\n  __repr__ = str_short\n\ndef typeof_rank1(r1):\n  col_s, row_s = jax.typeof(r1.col).sharding, jax.typeof(r1.row).sharding\n  sharding = col_s.update(spec=jax.P(col_s.spec[0], row_s.spec[0]))\n  return Rank1Ty((r1.col.shape[0], r1.row.shape[0]), sharding)\n\nregister_hitype(Rank1, typeof_rank1)\n```\n\nExample:\n```text\nclass Outer(VJPHiPrimitive):\n  def __init__(self, col_aval, row_aval):\n    if not (len(col_aval.shape) == 1 and len(row_aval.shape) == 1):\n      raise TypeError(f'bad outer factor types: {col_aval}, {row_aval}')\n    sharding = col_aval.sharding.update(\n        spec=jax.P(col_aval.sharding.spec[0], row_aval.sharding.spec[0]))\n    self.in_avals = (col_aval, row_aval)\n    self.out_aval = Rank1Ty((col_aval.shape[0], row_aval.shape[0]), sharding)\n    self.params = {}\n    super().__init__()\n\n  def expand(self, col, row):\n    return Rank1(col, row)\n\n  def vjp_fwd(self, nzs_in, col, row):\n    return self(col, row), (col, row)\n\n  def vjp_bwd_retval(self, res, g):   # g is dense, f32[m, n]\n    col, row = res\n    return g @ row, col @ g\n\nclass Factors(VJPHiPrimitive):\n  def __init__(self, r1_aval):\n    self.in_avals = (r1_aval,)\n    self.out_aval = tuple(r1_aval.lo_ty())  # the factor types are the lo types\n    self.params = {}\n    super().__init__()\n\n  def expand(self, r1):\n    return (r1.col, r1.row)\n\nclass MatmulR1(VJPHiPrimitive):\n  def __init__(self, x_aval, r1_aval):\n    if not (isinstance(r1_aval, Rank1Ty) and len(x_aval.shape) == 2 and\n            x_aval.shape[1] == r1_aval.shape[0]):\n      raise TypeError(f'bad matmul_r1 operand types: {x_aval} @ {r1_aval}')\n    self.in_avals = (x_aval, r1_aval)\n    out_sharding = x_aval.sharding.update(\n        spec=jax.P(x_aval.sharding.spec[0], r1_aval.sharding.spec[1]))\n    self.out_aval = ShapedArray((x_aval.shape[0], r1_aval.shape[1]),\n                                jnp.dtype('float32'), sharding=out_sharding)\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x, r1):\n    return jnp.outer(x @ r1.col, r1.row)   # never materialize col ⊗ row\n\n  def vjp_fwd(self, nzs_in, x, r1):\n    col, row = factors(r1)  # rules are traced code: primitives, not attributes\n    return self(x, r1), (x, col, row)\n\n  def vjp_bwd_retval(self, res, g):\n    x, col, row = res\n    return jnp.outer(g @ row, col), x.T @ g\n\ndef outer(col, row):\n  return Outer(jax.typeof(col), jax.typeof(row))(col, row)\n\ndef factors(r1):\n  return Factors(jax.typeof(r1))(r1)\n\ndef matmul_r1(x, r1):\n  return MatmulR1(jax.typeof(x), jax.typeof(r1))(x, r1)\n```\n\nExample:\n```text\ncol = jnp.arange(6., dtype='float32') / 6.\nrow = jnp.arange(5., dtype='float32') / 5.\nr1 = outer(col, row)\nprint(jax.typeof(r1))\n\nacts = jnp.ones((3, 6), 'float32')\nprint(jnp.max(jnp.abs(matmul_r1(acts, r1) - acts @ jnp.outer(col, row))))\n```\n\nExample:\n```text\nr1[6,5]\n2.3841858e-07\n```\n\nExample:\n```text\njax.jit(matmul_r1).trace(acts, r1).jaxpr\n```\n\nExample:\n```text\n{ lambda ; a:f32[3,6] b:r1[6,5]. let\n    c:f32[3,5] = call_hi_primitive[_prim=MatmulR1[{}]] a b\n  in (c,) }\n```\n\nExample:\n```text\ndef r1_loss(x, r1):\n  return jnp.sum(matmul_r1(x, r1) ** 2)\n\nprint(jax.typeof(jax.grad(r1_loss, argnums=1)(acts, r1)))\n\ndef factor_loss(col, row):\n  return jnp.sum(matmul_r1(acts, outer(col, row)) ** 2)\n\ng_col, g_row = jax.grad(factor_loss, argnums=(0, 1))(col, row)\nprint(jax.typeof(g_col), jax.typeof(g_row))\n```\n\nExample:\n```text\nfloat32[6,5]\nfloat32[6] float32[5]\n```\n\nExample:\n```text\nprint(jax.typeof(outer(jax.device_put(jnp.zeros(8, 'float32'), jax.P('i')),\n                       row)))\n```\n\nExample:\n```text\nr1[8@i,5]\n```\n\nExample:\n```text\nimport itertools\n\n@dataclass(frozen=True)\nclass HiTup:\n  elts: tuple\n\n@dataclass(frozen=True)\nclass TupTy(HiType):\n  tys: tuple  # component types: array types, or other hi types\n\n  # lowering delegates to the component types\n  def lo_ty(self):\n    return [lo for ty in self.tys for lo in ty.lo_ty()]\n  def lower_val(self, tup):\n    return [lo for ty, elt in zip(self.tys, tup.elts)\n            for lo in ty.lower_val(elt)]\n  def raise_val(self, *los):\n    los = iter(los)\n    return HiTup(tuple(ty.raise_val(*itertools.islice(los, len(ty.lo_ty())))\n                       for ty in self.tys))\n\n  # so does the tangent type\n  def to_tangent_aval(self):\n    return TupTy(tuple(ty.to_tangent_aval() for ty in self.tys))\n\n  def str_short(self, short_dtypes=False, mesh_axis_types=False):\n    return ('Tup{' +\n            ','.join(t.str_short(short_dtypes, mesh_axis_types)\n                     for t in self.tys) + '}')\n  __repr__ = str_short\n\nregister_hitype(HiTup, lambda t: TupTy(tuple(map(jax.typeof, t.elts))))\n```\n\nExample:\n```text\n@dataclass(frozen=True)\nclass TupSpec(MappingSpec):\n  val: tuple  # one axis entry per component\n\ndef tup_dec_rank(self, size, spec):\n  return TupTy(tuple(ty.dec_rank(size, s)\n                     for ty, s in zip(self.tys, spec.val)))\n\ndef tup_inc_rank(self, size, spec):\n  return TupTy(tuple(ty.inc_rank(size, s)\n                     for ty, s in zip(self.tys, spec.val)))\n\nTupTy.dec_rank = tup_dec_rank\nTupTy.inc_rank = tup_inc_rank\n\nclass MakeTup(VJPHiPrimitive):\n  def __init__(self, elt_avals):\n    self.in_avals = tuple(elt_avals)\n    self.out_aval = TupTy(tuple(elt_avals))\n    self.params = {}\n    super().__init__()\n\n  def expand(self, *elts):\n    return HiTup(elts)\n\n  def batch(self, axis_data, args, in_dims):\n    return make_tup(*args), TupSpec(tuple(in_dims))\n\nclass GetTupElt(VJPHiPrimitive):\n  def __init__(self, tup_aval, idx):\n    self.in_avals = (tup_aval,)\n    self.out_aval = tup_aval.tys[idx]\n    self.params = dict(idx=idx)\n    super().__init__()\n\n  def expand(self, tup):\n    return tup.elts[self.idx]\n\n  def batch(self, axis_data, args, in_dims):\n    tup, = args\n    spec, = in_dims\n    if spec is None:\n      return get_tuple_element(tup, self.idx), None\n    return get_tuple_element(tup, self.idx), spec.val[self.idx]\n\ndef make_tup(*elts):\n  return MakeTup(map(jax.typeof, elts))(*elts)\n\ndef get_tuple_element(tup, idx):\n  return GetTupElt(jax.typeof(tup), idx)(tup)\n```\n\nExample:\n```text\ntup = make_tup(jnp.arange(3.), 5.)\nprint(jax.typeof(tup))\nprint(get_tuple_element(tup, 1))\n\nnested = make_tup(make_tup(1., 2.), jnp.arange(2.))\nprint(jax.typeof(nested))\nprint(jax.jit(lambda t: get_tuple_element(get_tuple_element(t, 0), 1))(nested))\n\nprint(jax.typeof(make_tup(qx, 3.)))  # a quantized array element\n```\n\nExample:\n```text\nTup{float32[3],float32[]}\n5.0\nTup{Tup{float32[],float32[]},float32[2]}\n2.0\nTup{q8[2,3],float32[]}\n```\n\nExample:\n```text\ndef swap(t):\n  a, b = get_tuple_element(t, 0), get_tuple_element(t, 1)\n  return make_tup(b, a)\n\nout = jax.vmap(swap, in_axes=TupSpec((0, None)), out_axes=TupSpec((None, 0)),\n               axis_size=3)(tup)\nprint(jax.typeof(tup), '->', jax.typeof(out))\n```\n\nExample:\n```text\nTup{float32[3],float32[]} -> Tup{float32[],float32[3]}\n```\n\nExample:\n```text\nprint(jax.typeof(make_tup(rows, jnp.float32(1.))))\n```\n\nExample:\n```text\nTup{float32[8@i,3],float32[]}\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.694Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":76,"totalLines":901,"estimatedTokens":5405}}35{"id":"doc-jax_numpy_fft_rfftn_jax_documentation-2d4d3969","source":"documentation","title":"jax.numpy.fft.rfftn — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.rfftn.html","text":"Example:\n```text\n>>> x = jnp.array([[[1, 3, 5],\n...                 [2, 4, 6]],\n...                [[7, 9, 11],\n...                 [8, 10, 12]]])\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.rfftn(x)\nArray([[[ 78.+0.j  , -12.+6.93j],\n        [ -6.+0.j  ,   0.+0.j  ]],\n\n       [[-36.+0.j  ,   0.+0.j  ],\n        [  0.+0.j  ,   0.+0.j  ]]], dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.rfftn(x, s=[3, 3, 4])\nArray([[[ 78.   +0.j  , -16.  -26.j  ,  26.   +0.j  ],\n        [ 15.  -36.37j, -16.12 +1.93j,   5.  -12.12j],\n        [ 15.  +36.37j,   8.12-11.93j,   5.  +12.12j]],\n\n       [[ -7.5 -49.36j, -20.45 +9.43j,  -2.5 -16.45j],\n        [-25.5  -7.79j,  -0.6 +11.96j,  -8.5  -2.6j ],\n        [ 19.5 -12.99j,  -8.33 -6.5j ,   6.5  -4.33j]],\n\n       [[ -7.5 +49.36j,  12.45 -4.43j,  -2.5 +16.45j],\n        [ 19.5 +12.99j,   0.33 -6.5j ,   6.5  +4.33j],\n        [-25.5  +7.79j,   4.6  +5.04j,  -8.5  +2.6j ]]], dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.rfftn(x, s=[3, 5], axes=[0, 1])\nArray([[[ 18.   +0.j  ,  26.   +0.j  ,  34.   +0.j  ],\n        [ 11.09 -9.51j,  16.33-13.31j,  21.56-17.12j],\n        [ -0.09 -5.88j,   0.67 -8.23j,   1.44-10.58j]],\n\n       [[ -4.5 -12.99j,  -2.5 -16.45j,  -0.5 -19.92j],\n        [ -9.71 -6.3j , -10.05 -9.52j, -10.38-12.74j],\n        [ -4.95 +0.72j,  -5.78 -0.2j ,  -6.61 -1.12j]],\n\n       [[ -4.5 +12.99j,  -2.5 +16.45j,  -0.5 +19.92j],\n        [  3.47+10.11j,   6.43+11.42j,   9.38+12.74j],\n        [  3.19 +1.63j,   4.4  +1.38j,   5.61 +1.12j]]], dtype=complex64)\n```\n\nExample:\n```text\n>>> x1 = jnp.array([1, 2, 3, 4])\n>>> jnp.fft.rfftn(x1)\nArray([10.+0.j, -2.+2.j, -2.+0.j], dtype=complex64)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.699Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":57,"estimatedTokens":450}}36{"id":"doc-jax_array_committed_jax_documentation-54727424","source":"documentation","title":"jax.Array.committed — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.Array.committed.html","text":"Example:\n```text\n>>> a = jax.device_put(np.arange(8), jax.devices()[0])\n>>> b = jax.device_put(np.arange(8), jax.devices()[1])\n>>> a + b  \nTraceback (most recent call last):\n  ...\nValueError: Received incompatible devices for jitted computation.\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.700Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":1,"totalLines":11,"estimatedTokens":66}}37{"id":"doc-jax_numpy_module_jax_documentation-e7932e7d","source":"documentation","title":"jax.numpy module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.numpy.html","text":"Example:\n```text\n>>> def f(x):\n...   nx = x.__array_namespace__()\n...   return nx.sin(x) ** 2 + nx.cos(x) ** 2\n\n>>> import jax.numpy as jnp\n>>> x = jnp.arange(5)\n>>> f(x).round()\nArray([1., 1., 1., 1., 1.], dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.702Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":1,"totalLines":13,"estimatedTokens":60}}38{"id":"doc-jax_array_at_jax_documentation-2cbfa046","source":"documentation","title":"jax.Array.at — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.Array.at.html","text":"Example:\n```text\n>>> x = jnp.arange(5.0)\n>>> x\nArray([0., 1., 2., 3., 4.], dtype=float32)\n>>> x.at[2].get()\nArray(2., dtype=float32)\n>>> x.at[2].add(10)\nArray([ 0.,  1., 12.,  3.,  4.], dtype=float32)\n```\n\nExample:\n```text\n>>> x.at[10].add(10)  # dropped\nArray([0., 1., 2., 3., 4.], dtype=float32)\n>>> x.at[20].add(10, mode='clip')  # clipped\nArray([ 0.,  1.,  2.,  3., 14.], dtype=float32)\n```\n\nExample:\n```text\n>>> x.at[20].get()  # out-of-bounds indices clipped\nArray(4., dtype=float32)\n>>> x.at[20].get(mode='fill')  # out-of-bounds indices filled with NaN\nArray(nan, dtype=float32)\n>>> x.at[20].get(mode='fill', fill_value=-1)  # custom fill value\nArray(-1., dtype=float32)\n```\n\nExample:\n```text\n>>> x.at[-1].set(99)\nArray([ 0.,  1.,  2.,  3., 99.], dtype=float32)\n>>> x.at[-1].set(99, wrap_negative_indices=False, mode='drop')  # dropped!\nArray([0., 1., 2., 3., 4.], dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.705Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":38,"estimatedTokens":227}}39{"id":"doc-jax_array_byteswap_jax_documentation-45aa8b4e","source":"documentation","title":"jax.Array.byteswap — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.Array.byteswap.html","text":"Example:\n```text\n>>> import jax.numpy as jnp\n>>> x = jnp.arange(5, dtype='int32')\n>>> x\nArray([0, 1, 2, 3, 4], dtype=int32)\n>>> x.byteswap()\nArray([       0, 16777216, 33554432, 50331648, 67108864], dtype=int32)\n```\n\nExample:\n```text\n>>> import numpy as np\n>>> np.array(x.byteswap()).view('>i4')  # view as big-endian\narray([0, 1, 2, 3, 4], dtype='>i4')\n```\n\nExample:\n```text\n>>> x.byteswap().byteswap()\nArray([0, 1, 2, 3, 4], dtype=int32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.705Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":24,"estimatedTokens":115}}40{"id":"doc-jax_array_view_jax_documentation-4b78cc9b","source":"documentation","title":"jax.Array.view — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.Array.view.html","text":"Example:\n```text\n>>> jnp.zeros([1,2,3], dtype=jnp.int16).view(jnp.int8).shape\n(1, 2, 6)\n>>> jnp.zeros([1,2,4], dtype=jnp.int8).view(jnp.int16).shape\n(1, 2, 2)\n```\n\nExample:\n```text\n>>> jnp.array([1, 0, 1], dtype=jnp.int8).view(jnp.bool_)\nArray([ True, False,  True], dtype=bool)\n```\n\nExample:\n```text\n>>> jnp.array([1, 2, 0], dtype=jnp.int8) != 0\nArray([ True,  True, False], dtype=bool)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.719Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":21,"estimatedTokens":102}}41{"id":"doc-writing_tpu_kernels_with_pallas_jax_documentatio-cb5bf83f","source":"documentation","title":"Writing TPU kernels with Pallas — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/tpu/details.html","text":"Example:\n```text\npallas_call(\n    ...,\n    compiler_params=pltpu.CompilerParams(\n        dimension_semantics=[\"parallel\", \"parallel\", \"arbitrary\"]\n    ),\n  )\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.721Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":1,"totalLines":11,"estimatedTokens":44}}42{"id":"doc-quickstart_tpu_jax_documentation-ae1298e6","source":"documentation","title":"Quickstart: TPU — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/tpu/quickstart.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.experimental import pallas as pl\nfrom jax.experimental.pallas import tpu as pltpu\n```\n\nExample:\n```text\n@pl.kernel(\n    out_type=jax.ShapeDtypeStruct((128,), jnp.float32),\n    mesh=pltpu.TensorCoreMesh(axis_name='core'),\n    scratch_types=dict(o_vmem=pltpu.VMEM((128,), jnp.float32)),\n)\ndef fill_42(o_ref, o_vmem):\n  # Compute in VMEM\n  o_vmem[...] = jnp.full_like(o_vmem, 42.0)\n\n  # VMEM → HBM (blocks until the transfer completes)\n  pltpu.sync_copy(o_vmem, o_ref)\n\nresult = fill_42()  # [42.0, 42.0, ...]\n```\n\nExample:\n```text\ndef iota() -> jax.Array:\n  tpu_info = pltpu.get_tpu_info()\n\n  @pl.kernel(\n      out_type=jax.ShapeDtypeStruct((128 * tpu_info.num_cores,), jnp.float32),\n      mesh=pltpu.TensorCoreMesh(axis_name='core'),\n      scratch_types=dict(o_vmem=pltpu.VMEM((128,), jnp.float32)),\n  )\n  def kernel(o_ref, o_vmem):\n    i = jax.lax.axis_index('core')\n\n    # Compute our chunk in VMEM\n    o_vmem[...] = jnp.arange(128, dtype=jnp.float32) + i * 128\n\n    # Copy back to our slice of HBM\n    pltpu.sync_copy(o_vmem, o_ref.at[pl.ds(i * 128, 128)])\n\n  return kernel()\n\nresult = iota()  # [0.0, 1.0, 2.0, ...]\n```\n\nExample:\n```text\ndef add_matrices_pipelined(x: jax.Array, y: jax.Array) -> jax.Array:\n  @pl.kernel(\n      out_type=jax.ShapeDtypeStruct.like(x),\n      mesh=pltpu.TensorCoreMesh(axis_name='core'),\n  )\n  def kernel(x_hbm, y_hbm, o_hbm):\n    def add_body(x_vmem, y_vmem, o_vmem):\n      o_vmem[...] = x_vmem[...] + y_vmem[...]\n\n    pltpu.emit_pipeline(\n        add_body,\n        grid=(x_hbm.shape[0] // 128, x_hbm.shape[1] // 128),\n        in_specs=[\n            pl.BlockSpec((128, 128), lambda i, j: (i, j)),\n            pl.BlockSpec((128, 128), lambda i, j: (i, j)),\n        ],\n        out_specs=pl.BlockSpec((128, 128), lambda i, j: (i, j)),\n        core_axis_name='core',\n        dimension_semantics=(pltpu.PARALLEL, pltpu.ARBITRARY),\n    )(x_hbm, y_hbm, o_hbm)\n  return kernel(x, y)\n\nx = 2 * jnp.ones((256, 256))\ny = 3 * jnp.ones((256, 256))\nadd_matrices_pipelined(x, y)\n# Array([[5., 5., 5., ..., 5., 5., 5.],\n#        [5., 5., 5., ..., 5., 5., 5.],\n#        [5., 5., 5., ..., 5., 5., 5.],\n#        ...,\n#        [5., 5., 5., ..., 5., 5., 5.],\n#        [5., 5., 5., ..., 5., 5., 5.],\n#        [5., 5., 5., ..., 5., 5., 5.]], dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.721Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":86,"estimatedTokens":586}}43{"id":"doc-pseudo_random_number_generation_jax_documentatio-43fcc879","source":"documentation","title":"Pseudo-Random Number Generation — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/tpu/prng.html","text":"Example:\n```text\ndef body(key_ref, o_ref):\n  key = key_ref[...]\n  o_ref[...] = jax_random.uniform(\n      key, shape=o_ref[...].shape, minval=0.0, maxval=1.0\n  )\n\nthreefry_key = jax_random.key(0, impl=\"threefry2x32\")\n\n# We generate a threefry key outside of the kernel and pass it in via VMEM.\nresult = pl.pallas_call(\n    body,\n    in_specs=[pl.BlockSpec(memory_space=pltpu.VMEM)],\n    out_shape=jax.ShapeDtypeStruct((256, 256), jnp.float32)\n)(threefry_key)\n```\n\nExample:\n```text\nfrom jax.experimental.pallas import tpu as pltpu\n\ndef kernel_body(o_ref):\n  pltpu.prng_seed(0)\n  o_ref[...] = pltpu.stateful_uniform(shape=o_ref.shape, minval=0.0, maxval=1.0)\n\npl.pallas_call(kernel_body,\n               out_shape=jax.ShapeDtypeStruct((256, 256), jnp.float32))\n```\n\nExample:\n```text\ndef body(key_ref, o_ref):\n  o_ref[...] = jax.random.uniform(\n      key_ref[...], shape=o_ref[...].shape\n  )\n\nrbg_key = jax_random.key(0, impl=\"threefry2x32\")\nkey = pltpu.to_pallas_key(rbg_key)\no_shape = jax.ShapeDtypeStruct((8, 128), dtype)\nresult = pl.pallas_call(\n    body,\n    in_specs=[pl.BlockSpec(memory_space=pltpu.SMEM)],\n    out_shape=o_shape,\n)(key)\n```\n\nExample:\n```text\npltpu.sample_block(\n  sampler_function,  # A JAX random function, such as `jax.random.uniform`.\n  global_key,  # A global key shared across all blocks.\n  block_size,  # The local block size to generate.\n  tile_size,  # The tile size.\n  total_size,  # The total shape of the generated array across all blocks.\n  block_index,  # The block index into total_size. Usually this is the current program instance.\n  **sampler_kwargs  # Keyword arguments to sampler_function\n)\n```\n\nExample:\n```text\ndef make_kernel_body(index_map):\n  def body(key_ref, o_ref):\n    key = key_ref[...]\n    samples = pltpu.sample_block(\n        jax.random.uniform,\n        key,\n        block_size=o_ref[...].shape,\n        tile_size=(16, 128),\n        total_size=(64, 512),\n        block_index=index_map(pl.program_id(0), pl.program_id(1)),\n        minval=0.0,\n        maxval=1.0)\n    o_ref[...] = samples\n  return body\n\nglobal_key = pltpu.to_pallas_key(jax_random.key(0))\no_shape = jnp.ones((64, 512), dtype=jnp.float32)\nkey_spec = pl.BlockSpec(memory_space=pltpu.SMEM)\nout_spec = pl.BlockSpec((16, 128), lambda i, j: (i, j))\nresult_16x128 = pl.pallas_call(\n    make_kernel_body(index_map=lambda i, j: (i, j)),\n    out_shape=o_shape,\n    in_specs=[key_spec],\n    out_specs=out_spec,\n    grid=(4, 4),\n)(global_key)\n\nout_spec = pl.BlockSpec((32, 256), lambda i, j: (j, i))\nresult_32x256_transposed = pl.pallas_call(\n    make_kernel_body(index_map=lambda i, j: (j, i)),\n    in_specs=[key_spec],\n    out_shape=o_shape,\n    out_specs=out_spec,\n    grid=(2, 2),\n)(global_key)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.722Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":5,"totalLines":100,"estimatedTokens":680}}44{"id":"doc-quickstart_gpu_jax_documentation-54a40c00","source":"documentation","title":"Quickstart: GPU — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/gpu/quickstart.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.experimental import pallas as pl\nfrom jax.experimental.pallas import mosaic_gpu as plgpu\n```\n\nExample:\n```text\n@plgpu.kernel(out_type=jax.ShapeDtypeStruct((128,), jnp.float32))\ndef fill_42(o_ref):\n  o_ref[...] = jnp.full_like(o_ref, 42.0)\n\nresult = fill_42()  # [42.0, 42.0, ...]\n```\n\nExample:\n```text\n@plgpu.kernel(\n    out_type=jax.ShapeDtypeStruct((1024,), jnp.float32),\n    grid=(8,),\n    grid_names=('i',),\n)\ndef iota(o_ref):\n  i = jax.lax.axis_index('i')\n  o_ref[pl.ds(i * 128, 128)] = jnp.arange(128, dtype=jnp.float32) + i * 128\n\nresult = iota()  # [0.0, 1.0, ..., 1023.0]\n```\n\nExample:\n```text\ndef matmul(a, b, tile_m=128, tile_n=128, tile_k=64, out_dtype=jnp.float16):\n  m, k = a.shape\n  _, n = b.shape\n\n  @plgpu.kernel(\n      out_type=jax.ShapeDtypeStruct((m, n), out_dtype),\n      scratch_types=dict(\n          o_smem=plgpu.SMEM((tile_m, tile_n), out_dtype),\n          acc=plgpu.ACC((tile_m, tile_n), jnp.float32),\n      ),\n      grid=(m // tile_m, n // tile_n),\n      grid_names=('m', 'n'),\n  )\n  def kernel(a_gmem, b_gmem, o_gmem, o_smem, acc):\n    pid_m = jax.lax.axis_index('m')\n    pid_n = jax.lax.axis_index('n')\n\n    def body(_, a_smem, b_smem):\n      plgpu.wgmma(acc, a_smem, b_smem)\n      plgpu.wgmma_wait(1)  # Keep one wgmma in flight.\n\n    plgpu.emit_pipeline(\n        body,\n        grid=(k // tile_k,),\n        in_specs=[\n            plgpu.BlockSpec(\n                (tile_m, tile_k), lambda ki: (pid_m, ki), delay_release=1\n            ),\n            plgpu.BlockSpec(\n                (tile_k, tile_n), lambda ki: (ki, pid_n), delay_release=1\n            ),\n        ],\n        max_concurrent_steps=2,\n    )(a_gmem, b_gmem)\n\n    # Drain: move the accumulated result to GMEM via SMEM.\n    o_smem[...] = acc[...].astype(out_dtype)\n    plgpu.commit_smem()  # Make the SMEM write visible to the TMA engine.\n    plgpu.copy_smem_to_gmem(\n        o_smem,\n        o_gmem.at[pl.ds(pid_m * tile_m, tile_m),\n                  pl.ds(pid_n * tile_n, tile_n)],\n    )\n    plgpu.wait_smem_to_gmem(0)  # Wait for all copies to finish.\n\n  return kernel(a, b)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.722Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":82,"estimatedTokens":536}}45{"id":"doc-tpu_hardware_reference_jax_documentation-2424bebf","source":"documentation","title":"TPU Hardware Reference — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/tpu/hardware.html","text":"Example:\n```text\nfrom IPython.display import HTML, display\nfrom jax.experimental.pallas import tpu as pltpu\n\nheaders = [\n    \"Version\", \"Generation\", \"TensorCores/Chip\", \"VMEM Capacity\", \"CMEM Capacity\",\n    \"SMEM Capacity\", \"HBM Capacity\", \"HBM BW\", \"BF16 Peak\", \"FP8 Peak\", \"INT8 Peak\", \"INT4 Peak\", \"SparseCore\"\n]\n\nhtml_lines = []\nhtml_lines.append(\"<style>\")\nhtml_lines.append(\"  .bd-article {\")\nhtml_lines.append(\"    width: 1300px !important;\")\nhtml_lines.append(\"  }\")\nhtml_lines.append(\"  .tpu-spec-table {\")\nhtml_lines.append(\"    font-family: 'Google Sans', Arial, sans-serif;\")\nhtml_lines.append(\"    border-collapse: collapse;\")\nhtml_lines.append(\"    width: 100%;\")\nhtml_lines.append(\"    margin: 20px 0;\")\nhtml_lines.append(\"    font-size: 14px;\")\nhtml_lines.append(\"  }\")\nhtml_lines.append(\"  .tpu-spec-table th {\")\nhtml_lines.append(\"    background-color: #3c4043;\")\nhtml_lines.append(\"    color: white;\")\nhtml_lines.append(\"    text-align: left;\")\nhtml_lines.append(\"    padding: 12px 16px;\")\nhtml_lines.append(\"    font-weight: 500;\")\nhtml_lines.append(\"    border: 1px solid #dadce0;\")\nhtml_lines.append(\"  }\")\nhtml_lines.append(\"  .tpu-spec-table td {\")\nhtml_lines.append(\"    padding: 12px 16px;\")\nhtml_lines.append(\"    border: 1px solid #dadce0;\")\nhtml_lines.append(\"    color: #3c4043;\")\nhtml_lines.append(\"  }\")\nhtml_lines.append(\"  .tpu-spec-table tr:nth-child(even) {\")\nhtml_lines.append(\"    background-color: #f8f9fa;\")\nhtml_lines.append(\"  }\")\nhtml_lines.append(\"  .tpu-spec-table tr:hover {\")\nhtml_lines.append(\"    background-color: #f1f3f4;\")\nhtml_lines.append(\"  }\")\nhtml_lines.append(\"</style>\")\nhtml_lines.append(\"<table class='tpu-spec-table'>\")\nhtml_lines.append(\"  <thead>\")\nhtml_lines.append(\"    <tr>\")\nfor h in headers:\n  html_lines.append(f\"      <th>{h}</th>\")\nhtml_lines.append(\"    </tr>\")\nhtml_lines.append(\"  </thead>\")\nhtml_lines.append(\"  <tbody>\")\n\nfor cv in pltpu.ChipVersion:\n  if cv == pltpu.ChipVersion.TPU_7:\n    continue  # Skip TPU 7 as it is redundant with 7x\n\n  # Get per-TensorCore specs (num_cores=1)\n  info = pltpu.get_tpu_info_for_chip(cv, 1)\n\n  sc = info.sparse_core\n  sc_str = \"No\"\n  if sc is not None:\n    sc_str = f\"Yes ({sc.num_cores} SCs, {sc.num_subcores} subcores, {sc.vmem_capacity_bytes // 1024} KiB VMEM)\"\n\n  row = [\n      cv.value.upper(),\n      f\"TPU v{info.generation}\" if info.generation < 7 else f\"TPU {info.generation}\",\n      str(cv.num_physical_tensor_cores_per_chip),\n      f\"{info.vmem_capacity_bytes // (1024 * 1024)} MiB\",\n      f\"{info.cmem_capacity_bytes // 1024} KiB\" if info.cmem_capacity_bytes > 0 else \"N/A\",\n      f\"{info.smem_capacity_bytes // 1024} KiB\",\n      f\"{info.hbm_capacity_bytes // 1000000000} GB\",\n      f\"{info.mem_bw_bytes_per_second / 1e9:.1f} GB/s\" if info.mem_bw_bytes_per_second > 0 else \"N/A\",\n      f\"{int(round(info.bf16_ops_per_second / 1e12))} TFLOPs/s\" if info.bf16_ops_per_second > 0 else \"N/A\",\n      f\"{int(round(info.fp8_ops_per_second / 1e12))} TFLOPs/s\" if info.fp8_ops_per_second > 0 else \"N/A\",\n      f\"{int(round(info.int8_ops_per_second / 1e12))} TOPs/s\" if info.int8_ops_per_second > 0 else \"N/A\",\n      f\"{int(round(info.int4_ops_per_second / 1e12))} TOPs/s\" if info.int4_ops_per_second > 0 else \"N/A\",\n      sc_str,\n  ]\n\n  html_lines.append(\"    <tr>\")\n  for cell in row:\n    html_lines.append(f\"      <td>{cell}</td>\")\n  html_lines.append(\"    </tr>\")\n\nhtml_lines.append(\"  </tbody>\")\nhtml_lines.append(\"</table>\")\n\ndisplay(HTML(\"\\n\".join(html_lines)))\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.723Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":1,"totalLines":91,"estimatedTokens":876}}46{"id":"doc-collective_matrix_multiplication_jax_documentati-672fce95","source":"documentation","title":"Collective matrix multiplication — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/gpu/collective_matmul.html","text":"Example:\n```text\ndef all_gather_lhs_matmul(\n    lhs: jax.Array,\n    rhs: jax.Array,\n    axis_name,\n    *,\n    config: hopper_matmul_mgpu.TuningConfig,\n    dtype: jnp.dtype = jnp.bfloat16,\n) -> jax.Array:\n  if (num_devices := jax.device_count()) != jax.process_count():\n    raise ValueError(\"The kernel only supports one device per process\")\n  if (axis_size := lax.axis_size(axis_name)) != num_devices:\n    raise ValueError(\"The kernel can only work over all devices in a Mesh.\")\n  ...\n\n  m_shard, k = lhs.shape\n  _, n_shard = rhs.shape\n  tile_m, tile_n, tile_k = config.tile_m, config.tile_n, config.tile_k\n  cta_tile_m = tile_m * (1 + (config.wg_dimension == MatmulDimension.M))\n  num_sms = jax.extend.backend.get_default_device().core_count\n\n  def kernel_body(lhs_local_ref, rhs_ref, out_ref, scratch_ref):\n    ...\n\n  result, _ = plgpu.kernel(\n      kernel_body,\n      out_shape=[\n          # The output (with M gathered)\n          jax.ShapeDtypeStruct((axis_size * m_shard, n_shard), dtype),\n          # A scratch buffer for LHS all-gather\n          jax.ShapeDtypeStruct((axis_size - 1, m_shard, k), dtype),\n      ],\n      grid=(num_sms,),\n      num_threads=3, # The matmul kernel uses 3 threads: 2 compute and 1 memory\n      thread_name=\"wg\",\n  )(lhs, rhs)\n  return result\n```\n\nExample:\n```text\ndef all_gather_lhs_matmul(...):\n  def kernel_body(lhs_local_ref, rhs_ref, out_ref, scratch_ref, out_smem, received_sem):\n    wg_idx = lax.axis_index(\"wg\")\n    dev_id = lax.axis_index(axis_name)\n    # This device sends to dev_id - 1, forming a ring.\n    send_dev_id = lax.rem(dev_id + axis_size - 1, axis_size)\n    send_scratch_ref = plgpu.remote_ref(scratch_ref, send_dev_id)\n\n    def device_step(lhs_source_ref, device_offset):\n      # Invariant: lhs_source_ref contains A_{(dev_id + device_offset) % D}\n      # and is ready to be used for computation.\n\n      ...\n\n    # We peel the first step to read data directly from lhs_local_ref.\n    device_step(lhs_local_ref, 0)\n    @pl.loop(1, num_devices)\n    def _device_loop(device_offset):\n      device_step(scratch_ref.at[device_offset - 1], device_offset)\n```\n\nExample:\n```text\ndef all_gather_lhs_matmul(...):\n  ...\n\n  def kernel_body(lhs_local_ref, rhs_ref, out_ref, scratch_ref, out_smem, received_sem):\n    ...\n\n    def device_step(lhs_source_ref, device_offset):\n      # We are computing block (dev_id + device_offset) % D of the output.\n      out_device_idx = lax.rem(device_offset + dev_id, axis_size)\n      out_device_m_slice = pl.ds(out_device_idx * m_shard, m_shard)\n\n      # In step `device_offset`, we send A_{(dev_id + device_offset) % D} to\n      # the next device in the ring, into scratch slot `device_offset`.\n      # We also don't send on the last step since that would return the data\n      # back to its original source.\n      next_scratch_slot = device_offset\n      is_send_wg = wg_idx == 0 # Only one warpgroup per CTA sends\n      has_send_space = next_scratch_slot < axis_size - 1\n      should_send = is_send_wg & has_send_space\n\n      # This function will be called by hopper_matmul_mgpu.kernel in the body\n      # of its pipeline. We use it to take the tile of LHS loaded into SMEM and\n      # issue a TMA send to the next device in the ring.\n      def send_lhs(m_idx, n_idx, k_idx, a_smem, b_smem, send_ref, should_send):\n        del b_smem  # Unused.\n        # We only send when n_idx == 0 to avoid sending the same data\n        # multiple times when revisiting the left operand.\n        @pl.when(should_send & jnp.bool(n_idx == 0))\n        def _():\n          k_slice = pl.ds(k_idx * tile_k, tile_k)\n          m_slice = pl.ds(m_idx * cta_tile_m, cta_tile_m)\n          plgpu.copy_smem_to_gmem(a_smem, send_ref.at[m_slice, k_slice])\n          # Wait for previous copies to complete. We pass in delay_release=1\n          # to the pipeline in the matmul kernel to ensure that it doesn't\n          # overwrite the input until at least the next step completes, but it\n          # will not wait any longer.\n          plgpu.wait_smem_to_gmem(1, wait_read_only=True)\n\n      hopper_matmul_mgpu.kernel(\n          lhs_source_ref,  # LHS shard for this step\n          rhs_ref,  # RHS shard is always the same\n          out_ref.at[out_device_m_slice],  # Slice of output to update\n          out_smem,\n          config=config,\n          pipeline_callback=functools.partial(\n              send_lhs,\n              send_ref=send_scratch_ref.at[next_scratch_slot],\n              should_send=should_send,\n          ),\n          delay_release=1,\n      )\n\n      # Wait for the next scratch to arrive for the next step's computation.\n      # Each device signals its neighbor when it has finished sending.\n      @pl.when(should_send)\n      def _signal():\n        # Make sure our remote copy is done, then signal.\n        plgpu.wait_smem_to_gmem(0, wait_read_only=False)\n        pl.semaphore_signal(received_sem, device_id=send_dev_id)\n      @pl.when(has_send_space)\n      def _wait():\n        # Here, we wait for the data to arrive from the previous device in the\n        # ring. At each step, will expect to receive a signal from each SM.\n        # We use decrement=False to make this operation slightly faster, but\n        # this also means that we need to scale the expected number of signals\n        # by the number of steps taken so far (as the value only increases).\n        pl.semaphore_wait(received_sem, value=(device_offset + 1) * num_sms, decrement=False)\n\n    ...\n```\n\nExample:\n```text\nm_shard, n_shard, k = 1024, 1024, 1024\ndtype = jnp.float16\nmesh = jax.make_mesh((jax.device_count(),), (\"x\",),\n                     axis_types=(jax.sharding.AxisType.Explicit,))\nwith jax.set_mesh(mesh):\n  a = jax.random.normal(jax.random.key(1), (m_shard * jax.device_count(), k), dtype)\n  b = jax.random.normal(jax.random.key(2), (k, n_shard * jax.device_count()), dtype)\n  a = jax.sharding.reshard(a, P(\"x\", None))\n  b = jax.sharding.reshard(b, P(None, \"x\"))\n\n  # Example config for 8xH100. You might need to retune to your shape.\n  config = hopper_matmul_mgpu.TuningConfig(\n      tile_m=128, tile_n=128, tile_k=64, max_concurrent_steps=4,\n      grid_minor_dim=MatmulDimension.N, grid_tile_width=8,\n      wg_dimension=MatmulDimension.N,\n  )\n\n  kernel = jax.jit(\n      jax.shard_map(\n          functools.partial(all_gather_lhs_matmul, axis_name=\"x\", config=config),\n          out_specs=P(None, \"x\"),\n          check_vma=False,\n      )\n  )\n  c = kernel(a, b)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.724Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":166,"estimatedTokens":1606}}47{"id":"doc-grids_and_blockspecs_jax_documentation-7c308042","source":"documentation","title":"Grids and BlockSpecs — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/grid_blockspec.html","text":"Example:\n```text\npl.pallas_call(some_kernel, grid=(n,))(...)\n```\n\nExample:\n```text\nfor i in range(n):\n  some_kernel(...)\n```\n\nExample:\n```text\npl.pallas_call(some_kernel, grid=(n, m))(...)\n```\n\nExample:\n```text\nfor i in range(n):\n  for j in range(m):\n    some_kernel(...)\n```\n\nExample:\n```text\n>>> import jax\n>>> from jax.experimental import pallas as pl\n>>> def slices_for_invocation(x_shape: tuple[int, ...],\n...                           x_spec: pl.BlockSpec,\n...                           grid: tuple[int, ...],\n...                           invocation_indices: tuple[int, ...]) -> tuple[slice, ...]:\n...   assert len(invocation_indices) == len(grid)\n...   assert all(0 <= i < grid_size for i, grid_size in zip(invocation_indices, grid))\n...   block_indices = x_spec.index_map(*invocation_indices)\n...   assert len(x_shape) == len(x_spec.block_shape) == len(block_indices)\n...   elem_indices = []\n...   for x_size, block_size, block_idx in zip(x_shape, x_spec.block_shape, block_indices):\n...     start_idx = block_idx * block_size\n...     # At least one element of the block must be within bounds\n...     assert start_idx < x_size\n...     elem_indices.append(slice(start_idx, start_idx + block_size))\n...   return elem_indices\n```\n\nExample:\n```text\n>>> slices_for_invocation(x_shape=(100, 100),\n...                       x_spec = pl.BlockSpec((10, 20), lambda i, j: (i, j)),\n...                       grid = (10, 5),\n...                       invocation_indices = (2, 4))\n[slice(20, 30, None), slice(80, 100, None)]\n\n>>> # Same shape of the array and blocks, but we iterate over each block 4 times\n>>> slices_for_invocation(x_shape=(100, 100),\n...                       x_spec = pl.BlockSpec((10, 20), lambda i, j, k: (i, j)),\n...                       grid = (10, 5, 4),\n...                       invocation_indices = (2, 4, 0))\n[slice(20, 30, None), slice(80, 100, None)]\n\n>>> # An example when the block is partially out-of-bounds in the 2nd axis.\n>>> slices_for_invocation(x_shape=(100, 90),\n...                       x_spec = pl.BlockSpec((10, 20), lambda i, j: (i, j)),\n...                       grid = (10, 5),\n...                       invocation_indices = (2, 4))\n[slice(20, 30, None), slice(80, 100, None)]\n```\n\nExample:\n```text\n>>> def show_program_ids(x_shape, block_shape, grid,\n...                      index_map=lambda i, j: (i, j)):\n...   def program_ids_kernel(o_ref):  # Fill the output block with 10*program_id(1) + program_id(0)\n...     axes = 0\n...     for axis in range(len(grid)):\n...       axes += pl.program_id(axis) * 10**(len(grid) - 1 - axis)\n...     o_ref[...] = jnp.full(o_ref.shape, axes)\n...   res = pl.pallas_call(program_ids_kernel,\n...                        out_shape=jax.ShapeDtypeStruct(x_shape, dtype=np.int32),\n...                        grid=grid,\n...                        in_specs=[],\n...                        out_specs=pl.BlockSpec(block_shape, index_map),\n...                        interpret=True)()\n...   print(res)\n```\n\nExample:\n```text\n>>> show_program_ids(x_shape=(8, 6), block_shape=(2, 3), grid=(4, 2),\n...                  index_map=lambda i, j: (i, j))\n[[ 0  0  0  1  1  1]\n [ 0  0  0  1  1  1]\n [10 10 10 11 11 11]\n [10 10 10 11 11 11]\n [20 20 20 21 21 21]\n [20 20 20 21 21 21]\n [30 30 30 31 31 31]\n [30 30 30 31 31 31]]\n\n>>> # An example with out-of-bounds accesses\n>>> show_program_ids(x_shape=(7, 5), block_shape=(2, 3), grid=(4, 2),\n...                  index_map=lambda i, j: (i, j))\n[[ 0  0  0  1  1]\n [ 0  0  0  1  1]\n [10 10 10 11 11]\n [10 10 10 11 11]\n [20 20 20 21 21]\n [20 20 20 21 21]\n [30 30 30 31 31]]\n\n>>> # It is allowed for the shape to be smaller than block_shape\n>>> show_program_ids(x_shape=(1, 2), block_shape=(2, 3), grid=(1, 1),\n...                  index_map=lambda i, j: (i, j))\n[[0 0]]\n```\n\nExample:\n```text\n>>> show_program_ids(x_shape=(8, 6), block_shape=(2, 3), grid=(4, 2, 10),\n...                  index_map=lambda i, j, k: (i, j))\n[[  9   9   9  19  19  19]\n [  9   9   9  19  19  19]\n [109 109 109 119 119 119]\n [109 109 109 119 119 119]\n [209 209 209 219 219 219]\n [209 209 209 219 219 219]\n [309 309 309 319 319 319]\n [309 309 309 319 319 319]]\n```\n\nExample:\n```text\n>>> def kernel(o_ref):\n...   assert o_ref.shape == (2,)\n...   o_ref[...] = jnp.full((2,), 10 * pl.program_id(1) + pl.program_id(0))\n>>> pl.pallas_call(kernel,\n...                jax.ShapeDtypeStruct((3, 4), dtype=np.int32),\n...                out_specs=pl.BlockSpec((None, 2), lambda i, j: (i, j)),\n...                grid=(3, 2), interpret=True)()\nArray([[ 0,  0, 10, 10],\n       [ 1,  1, 11, 11],\n       [ 2,  2, 12, 12]], dtype=int32)\n```\n\nExample:\n```text\n>>> show_program_ids(x_shape=(4, 4), block_shape=None, grid=(2, 3),\n...                  index_map=None)\n[[12 12 12 12]\n [12 12 12 12]\n [12 12 12 12]\n [12 12 12 12]]\n\n>>> show_program_ids(x_shape=(4, 4), block_shape=(4, 4), grid=(2, 3),\n...                  index_map=None)\n[[12 12 12 12]\n [12 12 12 12]\n [12 12 12 12]\n [12 12 12 12]]\n```\n\nExample:\n```text\n>>> # element without padding\n>>> show_program_ids(x_shape=(8, 6), block_shape=(pl.Element(2), pl.Element(3)),\n...                  grid=(4, 2),\n...                  index_map=lambda i, j: (2*i, 3*j))\n    [[ 0  0  0  1  1  1]\n     [ 0  0  0  1  1  1]\n     [10 10 10 11 11 11]\n     [10 10 10 11 11 11]\n     [20 20 20 21 21 21]\n     [20 20 20 21 21 21]\n     [30 30 30 31 31 31]\n     [30 30 30 31 31 31]]\n\n>>> # element, first pad the array with 1 row and 2 columns.\n>>> show_program_ids(x_shape=(7, 7),\n...                  block_shape=(pl.Element(2, (1, 0)),\n...                               pl.Element(3, (2, 0))),\n...                  grid=(4, 3),\n...                  index_map=lambda i, j: (2*i, 3*j))\n    [[ 0  1  1  1  2  2  2]\n     [10 11 11 11 12 12 12]\n     [10 11 11 11 12 12 12]\n     [20 21 21 21 22 22 22]\n     [20 21 21 21 22 22 22]\n     [30 31 31 31 32 32 32]\n     [30 31 31 31 32 32 32]]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.725Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":12,"totalLines":191,"estimatedTokens":1476}}48{"id":"doc-pallas_quickstart_jax_documentation-8d21b0ed","source":"documentation","title":"Pallas Quickstart — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/quickstart.html","text":"Example:\n```text\nfrom functools import partial\n\nimport jax\nfrom jax.experimental import pallas as pl\nimport jax.numpy as jnp\nimport numpy as np\n```\n\nExample:\n```text\ndef add_vectors_kernel(x_ref, y_ref, o_ref):\n  x, y = x_ref[...], y_ref[...]\n  o_ref[...] = x + y\n```\n\nExample:\n```text\ndef add_sliced_kernel(x_ref, y_ref, o_ref):\n  small_mid = x_ref.shape[0] // 2\n\n  x_left = x_ref.at[:small_mid]\n  x_right = x_ref.at[small_mid:]\n  y_left = y_ref.at[:small_mid]\n  y_right = y_ref.at[small_mid:]\n\n  # The output shape is (4*small_mid).\n  large_mid = 2*small_mid\n  o_ref.at[:large_mid][:small_mid] = x_left[...] + y_left[...]\n  o_ref.at[:large_mid][small_mid:] = x_left[...] + y_right[...]\n  o_ref.at[large_mid:][:small_mid] = x_right[...] + y_left[...]\n  o_ref.at[large_mid:][small_mid:] = x_right[...] + y_right[...]\n```\n\nExample:\n```text\n@jax.jit\ndef add_vectors(x: jax.Array, y: jax.Array) -> jax.Array:\n  return pl.pallas_call(\n      add_vectors_kernel,\n      out_shape=jax.ShapeDtypeStruct.like(x)\n  )(x, y)\nadd_vectors(jnp.arange(8), jnp.arange(8))\n```\n\nExample:\n```text\nArray([ 0,  2,  4,  6,  8, 10, 12, 14], dtype=int32)\n```\n\nExample:\n```text\ndef iota_kernel(o_ref):\n  i = pl.program_id(0)\n  o_ref[i] = i\n```\n\nExample:\n```text\n# GPU version\ndef iota(size: int):\n  return pl.pallas_call(iota_kernel,\n                        out_shape=jax.ShapeDtypeStruct((size,), jnp.int32),\n                        grid=(size,))()\niota(8)\n```\n\nExample:\n```text\nArray([0, 1, 2, 3, 4, 5, 6, 7], dtype=int32)\n```\n\nExample:\n```text\n# TPU version\nfrom jax.experimental.pallas import tpu as pltpu\n\ndef iota(size: int):\n  return pl.pallas_call(iota_kernel,\n                        out_specs=pl.BlockSpec(memory_space=pltpu.SMEM),\n                        out_shape=jax.ShapeDtypeStruct((size,), jnp.int32),\n                        grid=(size,))()\niota(8)\n```\n\nExample:\n```text\ndef matmul_kernel(x_ref, y_ref, z_ref):\n  z_ref[...] = x_ref[...] @ y_ref[...]\n\ndef matmul(x: jax.Array, y: jax.Array):\n  return pl.pallas_call(\n    matmul_kernel,\n    out_shape=jax.ShapeDtypeStruct((x.shape[0], y.shape[1]), x.dtype),\n    grid=(2, 2),\n    in_specs=[\n        pl.BlockSpec((x.shape[0] // 2, x.shape[1]), lambda i, j: (i, 0)),\n        pl.BlockSpec((y.shape[0], y.shape[1] // 2), lambda i, j: (0, j))\n    ],\n    out_specs=pl.BlockSpec(\n        (x.shape[0] // 2, y.shape[1] // 2), lambda i, j: (i, j),\n    )\n  )(x, y)\nk1, k2 = jax.random.split(jax.random.key(0))\nx = jax.random.normal(k1, (1024, 1024))\ny = jax.random.normal(k2, (1024, 1024))\nz = matmul(x, y)\nnp.testing.assert_allclose(z, x @ y)\n```\n\nExample:\n```text\ndef matmul_kernel(x_ref, y_ref, z_ref, *, activation):\n  z_ref[...] = activation(x_ref[...] @ y_ref[...])\n\ndef matmul(x: jax.Array, y: jax.Array, *, activation):\n  return pl.pallas_call(\n    partial(matmul_kernel, activation=activation),\n    out_shape=jax.ShapeDtypeStruct((x.shape[0], y.shape[1]), x.dtype),\n    grid=(2, 2),\n    in_specs=[\n        pl.BlockSpec((x.shape[0] // 2, x.shape[1]), lambda i, j: (i, 0)),\n        pl.BlockSpec((y.shape[0], y.shape[1] // 2), lambda i, j: (0, j))\n    ],\n    out_specs=pl.BlockSpec(\n        (x.shape[0] // 2, y.shape[1] // 2), lambda i, j: (i, j)\n    ),\n  )(x, y)\nk1, k2 = jax.random.split(jax.random.key(0))\nx = jax.random.normal(k1, (1024, 1024))\ny = jax.random.normal(k2, (1024, 1024))\nz = matmul(x, y, activation=jax.nn.relu)\nnp.testing.assert_allclose(z, jax.nn.relu(x @ y))\n```\n\nExample:\n```text\nk1, k2 = jax.random.split(jax.random.key(0))\nx = jax.random.normal(k1, (4, 1024, 1024))\ny = jax.random.normal(k2, (4, 1024, 1024))\nz = jax.vmap(partial(matmul, activation=jax.nn.relu))(x, y)\nnp.testing.assert_allclose(z, jax.nn.relu(jax.vmap(jnp.matmul)(x, y)))\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.726Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":12,"totalLines":146,"estimatedTokens":929}}49{"id":"doc-tpu_pipelining_jax_documentation-cde87a17","source":"documentation","title":"TPU Pipelining — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/tpu/pipelining.html","text":"Example:\n```text\n#@title Imports\n\nimport jax\nfrom jax.experimental import pallas as pl\nfrom jax.experimental.pallas import tpu as pltpu\nimport jax.numpy as jnp\nimport numpy as np\n```\n\nExample:\n```text\ndef hbm_vmem_kernel(x_hbm_ref, out_vmem_ref, scratch_vmem_ref):\n  pltpu.sync_copy(x_hbm_ref.at[0:1], scratch_vmem_ref)\n  out_vmem_ref[...] = scratch_vmem_ref[...] + 1\n\nx = jax.random.uniform(jax.random.key(0), (8, 128), jnp.float32)\nout = pl.pallas_call(hbm_vmem_kernel,\n  in_specs=[pl.BlockSpec(memory_space=pl.ANY)],\n  out_shape=jax.ShapeDtypeStruct((1, 128), jnp.float32),\n  scratch_shapes=(pltpu.VMEM(shape=(1, 128), dtype=jnp.float32),)\n)(x)\n\nnp.testing.assert_allclose(out, x[0:1] + 1)\n```\n\nExample:\n```text\npl.BlockSpec(\n  pipeline_mode=pl.Buffered(buffer_count=buffer_count)\n)\n```\n\nExample:\n```text\ndef emit_pipeline(\n    kernel: Callable,\n    grid: tuple[int],\n    in_specs: PyTree[BlockSpec] = None,\n    out_specs: PyTree[BlockSpec] = None,\n    dimension_semantics: tuple[GridDimensionSemantics] = None,\n    core_axis: int | None = None,\n) -> Callable:\n  ... # Returns a custom pipeline given an inner kernel and BlockSpecs.\n```\n\nExample:\n```text\npl.BlockSpec(\n  pipeline_mode=pl.Buffered(buffer_count=buffer_count, use_lookahead=True)\n)\n```\n\nExample:\n```text\npl.BlockSpec(\n   block_shape=(pl.BoundedSlice(32), 256),\n   index_map=lambda *grid_idxs: (pl.ds(start, end), 0),\n)\n```\n\nExample:\n```text\n# The following kernel copies `x` to the output in dynamic-sized chunks\n# passed in via `slices`.\n\ndef dynamic_block_example_kernel(x_hbm, slices_hbm, o_hbm, slices_smem):\n    pltpu.sync_copy(slices_hbm, slices_smem)  # Copy slices into SMEM.\n    def pipeline_body(x_vmem, o_vmem):\n        o_vmem[...] = x_vmem[...]\n    def index_map(i):\n        start = slices_smem[i, 0]\n        size = slices_smem[i, 1] - slices_smem[i, 0]\n        return (pl.ds(start, size), 0)\n    block_spec = pl.BlockSpec(block_shape=(pl.BoundedSlice(8), 128),\n                              index_map=index_map)\n    pltpu.emit_pipeline(\n        pipeline_body,\n        grid=(slices.shape[0],),\n        in_specs=[block_spec],\n        out_specs=block_spec\n    )(x_hbm, o_hbm)\n\nx = jax.random.uniform(jax.random.key(0), (8, 128), jnp.float32)\nslices = jnp.array([[0, 2], [2, 3], [3, 5], [5, 8]], dtype=jnp.int32)\n\nhbm_block_spec = pl.BlockSpec(memory_space=pl.ANY)\nout = pl.pallas_call(dynamic_block_example_kernel,\n               in_specs=[hbm_block_spec, hbm_block_spec],\n               out_specs=hbm_block_spec,\n               out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),\n               scratch_shapes=(pltpu.SMEM(slices.shape, jnp.int32),)\n              )(x, slices)\n\nnp.testing.assert_allclose(x, out)\n```\n\nExample:\n```text\ndef add_matrices_kernel(x_vmem_ref, y_vmem_ref, z_vmem_ref):\n  # Load x and y from VMEM into VREGs\n  x_vregs = x_vmem_ref[:, :]\n  y_vregs = y_vmem_ref[:, :]\n  # Execute a vectorized add\n  z_vregs = x_vregs + y_vregs\n  # Store the output values in VREGs back into VMEM\n  z_vmem_ref[:, :] = z_vregs\n\ndef add_matrices_pipelined_megacore(x: jax.Array, y: jax.Array) -> jax.Array:\n  block_spec = pl.BlockSpec((256, 512), lambda i: (i, 0))\n  return pl.pallas_call(\n      add_matrices_kernel,\n      out_shape=jax.ShapeDtypeStruct.like(x),\n      in_specs=[block_spec, block_spec],\n      out_specs=block_spec,\n      grid=(2,),\n      compiler_params=pltpu.CompilerParams(\n          dimension_semantics=(\"parallel\",))\n  )(x, y)\n\nx, y = jnp.ones((512, 512)), jnp.ones((512, 512))\nadd_matrices_pipelined_megacore(x, y)\n```\n\nExample:\n```text\nArray([[2., 2., 2., ..., 2., 2., 2.],\n       [2., 2., 2., ..., 2., 2., 2.],\n       [2., 2., 2., ..., 2., 2., 2.],\n       ...,\n       [2., 2., 2., ..., 2., 2., 2.],\n       [2., 2., 2., ..., 2., 2., 2.],\n       [2., 2., 2., ..., 2., 2., 2.]], dtype=float32)\n```\n\nExample:\n```text\ndef kernel_body(...):\n  def inner_pipeline_body(...):\n    ...\n  pltpu.emit_pipeline(inner_pipeline_body,\n                      grid=(4, 4), \n                      core_axis=0,\n                      dimension_semantics=(\"parallel\", \"sequential\"))\n\npl.pallas_call(\n      kernel_body,\n      grid=(num_cores,),\n      compiler_params=pltpu.CompilerParams(\n          dimension_semantics=(\"parallel\",))\n  )\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.727Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":10,"totalLines":155,"estimatedTokens":1060}}50{"id":"doc-mosaic_gpu_pipelining_jax_documentation-f0884c43","source":"documentation","title":"Mosaic GPU Pipelining — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/gpu/pipelining.html","text":"Example:\n```text\nimport jax\nfrom jax import lax\nfrom jax import numpy as jnp\nfrom jax.experimental.pallas import mosaic_gpu as plgpu\nfrom jax.experimental import pallas as pl\nimport numpy as np\n```\n\nExample:\n```text\ndef matmul(a, b, tile_m=128, tile_n=128, swizzle=128):\n  dtype = jnp.float16\n  swizzle_elems = swizzle // jnp.dtype(dtype).itemsize\n  tile_k = swizzle_elems\n  grid_m = m // tile_m\n  grid_k = k // tile_k\n  grid_n = n // tile_n\n  assert tile_m % swizzle_elems == 0\n\n  # Note: Transforms will be inferred automatically\n  # by Mosaic GPU in the future.\n  transforms = (\n    plgpu.TilingTransform((8, swizzle_elems)),\n    plgpu.SwizzleTransform(swizzle),\n  )\n\n  def kernel(a_gmem, b_gmem, o_gmem, o_smem, acc):\n    def pipeline_step(_, a_smem, b_smem):\n      plgpu.wgmma(acc, a_smem, b_smem)\n      plgpu.wgmma_wait(1)\n\n    # pl.program_id obtains the index into the grid.\n    pid_m = pl.program_id(0)\n    pid_n = pl.program_id(1)\n\n    pipeline = plgpu.emit_pipeline(\n        pipeline_step,\n        in_specs=[\n            plgpu.BlockSpec(\n                (tile_m, tile_k), lambda k: (pid_m, k), transforms=transforms\n            ),\n            plgpu.BlockSpec(\n                (tile_k, tile_n), lambda k: (k, pid_n), transforms=transforms\n            ),\n        ],\n        grid=(grid_k,),\n        max_concurrent_steps=2,\n        delay_release=1,\n    )\n\n    pipeline(a_gmem, b_gmem)\n    # Store WGMMA accumulator to SMEM and then to GMEM.\n    o_smem[...] = acc[...].astype(dtype)\n    plgpu.commit_smem()\n    m_slice = pl.ds(pid_m * tile_m, tile_m)\n    n_slice = pl.ds(pid_n * tile_n, tile_n)\n    plgpu.copy_smem_to_gmem(o_smem, o_gmem.at[m_slice, n_slice])\n    plgpu.wait_smem_to_gmem(0)\n\n  return plgpu.kernel(\n      kernel,\n      out_shape=jax.ShapeDtypeStruct((m, n), jnp.float16),\n      scratch_shapes=dict(\n          o_smem=plgpu.SMEM((tile_m, tile_n), jnp.float16),\n          acc=plgpu.ACC((tile_m, tile_n), jnp.float32)\n      ),\n      # grid specifies the CUDA grid.\n      # Instances of `kernel` will be executed in parallel over this grid.\n      grid=(grid_m, grid_n),\n      grid_names=(\"m\", \"n\"),\n  )(a, b)\n\nm = 132 * 128\nn = 4 * 128\nk = 10 * 64\nkey1, key2 = jax.random.split(jax.random.key(42), 2)\na = jax.random.uniform(key1, shape=(m, k), dtype=jnp.float16)\nb = jax.random.uniform(key2, shape=(k, n), dtype=jnp.float16)\n\nresult = matmul(a, b)\n\nnp.testing.assert_allclose(result, a @ b)\n```\n\nExample:\n```text\nplgpu.emit_pipeline_warp_specialized(\n  body: Callable,\n  *\n  grid: tuple[int, ...],\n  in_specs: Sequence[pallas_core.BlockSpec] = (),\n  out_specs: Sequence[pallas_core.BlockSpec] = (),\n  max_concurrent_steps: int,\n  compute_context: Callable\n  num_compute_wgs: int,\n  memory_registers: int\n  wg_axis: str,\n  memory_thread_idx: int | None = None,\n)\n```\n\nExample:\n```text\ndef matmul_warp_specialized(a, b, tile_m=128, tile_n=128, swizzle=128,\n                            compute_wgs=2):\n  dtype = jnp.float16\n  elems_128b = swizzle // jnp.dtype(dtype).itemsize\n  tile_k = elems_128b\n  grid_m = m // tile_m\n  grid_k = k // tile_k\n  grid_n = n // tile_n\n  assert tile_m % elems_128b == 0\n\n  transforms = (\n          plgpu.TilingTransform((8, elems_128b)),\n          plgpu.SwizzleTransform(128),\n      )\n\n  def kernel(a_gmem, b_gmem, o_gmem, o_smem):\n    wg_idx = lax.axis_index(\"wg\")\n    wg_slice = pl.ds(wg_idx * tile_n, tile_n)\n    # pl.program_id obtains the index into the pallas_call grid.\n    pid_m = pl.program_id(0)\n    pid_n = pl.program_id(1)\n\n    def compute_thread(pipeline):\n      acc = plgpu.layout_cast(\n          jnp.full((tile_m, tile_n), 0, dtype=jnp.float32), plgpu.Layout.WGMMA,\n      )\n      # yield marks the place where the pipelined loop will be inserted.\n      # Its argument are the initial carry values, and its result is the carry\n      # value after the loop completes.\n      final_acc = pipeline(acc)\n      o_smem[:, wg_slice] = final_acc[...].astype(dtype)\n\n    def kernel_body(_, a_smem, b_smem, carry):\n      acc = carry\n      b_smem_wg = b_smem.at[:, wg_slice]\n      def do_wgmma(acc_ref):\n        plgpu.wgmma(acc_ref, a_smem, b_smem_wg)\n      acc = pl.run_state(do_wgmma)(\n                          plgpu.ACC.init(acc))\n      return acc\n\n    pipeline = plgpu.emit_pipeline_warp_specialized(\n        kernel_body,\n        in_specs=[\n            plgpu.BlockSpec(\n              (tile_m, tile_k), lambda k: (pid_m, k), transforms=transforms\n            ),\n            plgpu.BlockSpec(\n              (tile_k, tile_n * 2), lambda k: (k, pid_n),transforms=transforms\n            ),\n        ],\n        grid=(grid_k,),\n        compute_context=compute_thread,\n        max_concurrent_steps=2,\n        num_compute_wgs=compute_wgs,\n        memory_registers=40,\n        memory_thread_idx=2,\n        wg_axis=\"wg\",\n    )\n    # Call the pipeline\n    pipeline(a_gmem, b_gmem)\n    # Copy the output from SMEM to GMEM.\n    plgpu.commit_smem()\n    m_slice = pl.ds(pid_m * tile_m, tile_m)\n    n_slice = pl.ds(pid_n * tile_n * 2, tile_n * 2)\n    plgpu.copy_smem_to_gmem(o_smem, o_gmem.at[m_slice, n_slice])\n    plgpu.wait_smem_to_gmem(0)\n\n  return plgpu.kernel(\n      kernel,\n      out_shape=jax.ShapeDtypeStruct((m, n), jnp.float16),\n      scratch_shapes=dict(\n          o_smem=plgpu.SMEM((tile_m, tile_n * 2), jnp.float16)\n      ),\n      grid=(grid_m, grid_n // 2),\n      grid_names=(\"m\", \"n\"),\n      num_threads=3,  # 2 compute, 1 memory.\n      thread_name=\"wg\"\n  )(a, b)\n\nm = 132 * 128\nn = 4 * 128\nk = 10 * 64\nkey1, key2 = jax.random.split(jax.random.key(42), 2)\na = jax.random.uniform(key1, shape=(m, k), dtype=jnp.float16)\nb = jax.random.uniform(key2, shape=(k, n), dtype=jnp.float16)\n\nresult = matmul_warp_specialized(a, b)\n\nnp.testing.assert_allclose(result, a @ b)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.728Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":198,"estimatedTokens":1433}}51{"id":"doc-matrix_multiplication_jax_documentation-8255a9fe","source":"documentation","title":"Matrix Multiplication — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/tpu/matmul.html","text":"Example:\n```text\n#@title Imports\nimport functools\nfrom typing import Callable\n\nimport jax\nfrom jax.experimental import pallas as pl\nfrom jax.experimental.pallas import tpu as pltpu\nfrom jax import random\nimport jax.numpy as jnp\nimport numpy as np\n```\n\nExample:\n```text\ndef matmul_small(x: np.ndarray, y: np.ndarray) -> np.ndarray:\n  m, k, n = x.shape[0], x.shape[1], y.shape[0]\n  assert m <= 256\n  assert k <= 256\n  assert n <= 256\n  return np.matmul(x, y)\n\ndef block_matmul(\n    x: np.ndarray,\n    y: np.ndarray,\n    *,\n    bm: int = 256,\n    bk: int = 256,\n    bn: int = 256,\n) -> np.ndarray:\n  m, k = x.shape\n  _, n = y.shape\n\n  z = np.zeros((m, n), dtype=x.dtype)\n  for m_i in range(m // bm):\n    for n_i in range(n // bn):\n      for k_i in range(k // bk):\n        m_slice = slice(m_i * bm, (m_i + 1) * bm)\n        k_slice = slice(k_i * bk, (k_i + 1) * bk)\n        n_slice = slice(n_i * bn, (n_i + 1) * bn)\n        x_block = x[m_slice, k_slice]\n        y_block = y[k_slice, n_slice]\n        z[m_slice, n_slice] += matmul_small(x_block, y_block)\n  return z\n```\n\nExample:\n```text\nm, k, n = 4096, 4096, 4096\nx = np.random.uniform(size=(m, k)).astype(np.float32)\ny = np.random.uniform(size=(k, n)).astype(np.float32)\nnp.testing.assert_allclose(x @ y, block_matmul(x, y), atol=1e-6, rtol=1e-6)\n```\n\nExample:\n```text\ndef matmul_kernel(x_ref, y_ref, z_ref):\n  @pl.when(pl.program_id(2) == 0)\n  def _():\n    z_ref[...] = jnp.zeros_like(z_ref)\n\n  z_ref[...] += x_ref[...] @ y_ref[...]\n\ndef matmul(\n    x: jax.Array,\n    y: jax.Array,\n    *,\n    bm: int = 128,\n    bk: int = 128,\n    bn: int = 128,\n):\n  m, k = x.shape\n  _, n = y.shape\n  return pl.pallas_call(\n      matmul_kernel,\n      out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),\n      in_specs=[pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)),\n                pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))],\n      out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),\n      grid=(m // bm, n // bn, k // bk),\n      compiler_params=pltpu.CompilerParams(\n          dimension_semantics=(\"parallel\", \"parallel\", \"arbitrary\")),\n  )(x, y)\n```\n\nExample:\n```text\nm, k, n = 4096, 4096, 4096\nk1, k2 = random.split(random.key(0), 2)\nx = random.normal(k1, (m, k), dtype=jnp.float32)\ny = random.normal(k2, (k, n), dtype=jnp.float32)\nnp.testing.assert_array_equal(x @ y, matmul(x, y))\n```\n\nExample:\n```text\ndef matmul_flops(m: int, k: int, n: int):\n  return 2 * m * k * n\n\ndef matmul_membw(m: int, k: int, n: int, dtype: jnp.dtype):\n  return (m * k + k * n + m * n) * np.dtype(dtype).itemsize\n\nprint(matmul_flops(1024, 1024, 1024))\nprint(matmul_membw(1024, 1024, 1024, jnp.float32))\n```\n\nExample:\n```text\n2147483648\n12582912\n```\n\nExample:\n```text\nv5e_flops = 197e12\nv5e_membw = 819e9\nv5e_op_intensity = v5e_flops / v5e_membw  # ~240.5\n```\n\nExample:\n```text\ndef matmul_flops_intensity(m: int, k: int, n: int, dtype: jnp.dtype):\n  flops = matmul_flops(m, k, n)\n  membw = matmul_membw(m, k, n, dtype)\n  return flops / membw\n```\n\nExample:\n```text\nprint(f\"{matmul_flops_intensity(1024, 1024, 1024, jnp.float32)} flops/byte\")\n```\n\nExample:\n```text\n170.66666666666666 flops/byte\n```\n\nExample:\n```text\nprint(f\"{matmul_flops_intensity(1024, 1024, 1024, jnp.bfloat16)} flops/byte\")\n```\n\nExample:\n```text\n341.3333333333333 flops/byte\n```\n\nExample:\n```text\ndef matmul_kernel(x_ref, y_ref, z_ref, acc_ref, *, nsteps):\n  @pl.when(pl.program_id(2) == 0)\n  def _():\n    acc_ref[...] = jnp.zeros_like(acc_ref)\n\n  acc_ref[...] += jnp.dot(\n      x_ref[...], y_ref[...], preferred_element_type=jnp.float32\n  )\n\n  @pl.when(pl.program_id(2) == nsteps - 1)\n  def _():\n    z_ref[...] = acc_ref[...].astype(z_ref.dtype)\n\n\n@jax.jit(static_argnames=['bm', 'bk', 'bn'])\ndef matmul(\n    x: jax.Array,\n    y: jax.Array,\n    *,\n    bm: int = 128,\n    bk: int = 128,\n    bn: int = 128,\n):\n  m, k = x.shape\n  _, n = y.shape\n  return pl.pallas_call(\n      functools.partial(matmul_kernel, nsteps=k // bk),\n      grid_spec=pltpu.PrefetchScalarGridSpec(\n        num_scalar_prefetch=0,\n        in_specs=[\n            pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)),\n            pl.BlockSpec((bk, bn), lambda i, j, k: (k, j)),\n        ],\n        out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),\n        scratch_shapes=[pltpu.VMEM((bm, bn), jnp.float32)],\n        grid=(m // bm, n // bn, k // bk),\n      ),\n      out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),\n      compiler_params=pltpu.CompilerParams(\n          dimension_semantics=(\"parallel\", \"parallel\", \"arbitrary\")),\n  )(x, y)\n```\n\nExample:\n```text\nm, k, n = 4096, 4096, 4096\nk1, k2 = random.split(random.key(0), 2)\nx = random.normal(k1, (m, k), dtype=jnp.bfloat16)\ny = random.normal(k2, (k, n), dtype=jnp.bfloat16)\nnp.testing.assert_array_equal(x @ y, matmul(x, y))\n```\n\nExample:\n```text\nimport timeit\n\ndef benchmark(f, ntrials: int = 100):\n  def run(*args, **kwargs):\n    # Compile function first\n    jax.block_until_ready(f(*args, **kwargs))\n    # Time function\n    result = timeit.timeit(lambda: jax.block_until_ready(f(*args, **kwargs)),\n                           number=ntrials)\n    time = result / ntrials\n    # print(f\"Time: {time}\")\n    return time\n  return run\n\ndef analyze_matmul(m: int, k: int, n: int, dtype: np.dtype,\n                   mm_func):\n  x = jnp.ones((m, k), dtype=dtype)\n  y = jnp.ones((k, n), dtype=dtype)\n  time = benchmark(mm_func)(x, y)\n  print(f\"----- {m} x {k} x {n} -----\")\n  print(\"Matmul time: \", time)\n  mm_flops = matmul_flops(m, k, n) / time\n  print(\"Matmul FLOP/s: \", mm_flops)\n  print(f\"FLOP/s utilization: {mm_flops / v5e_flops * 100:.4f}%\")\n  print()\n\nprint(\"================bm=128, bk=128, bn=128===================\")\nmm = functools.partial(matmul, bm=128, bk=128, bn=128)\nanalyze_matmul(1024, 1024, 1024, jnp.bfloat16, mm)\nanalyze_matmul(4096, 4096, 4096, jnp.bfloat16, mm)\nanalyze_matmul(8192, 8192, 8192, jnp.bfloat16, mm)\n\nprint(\"================bm=512, bk=1024, bn=1024===================\")\nmm = functools.partial(matmul, bm=512, bk=1024, bn=1024)\nanalyze_matmul(1024, 1024, 1024, jnp.bfloat16, mm)\nanalyze_matmul(4096, 4096, 4096, jnp.bfloat16, mm)\nanalyze_matmul(8192, 8192, 8192, jnp.bfloat16, mm)\n```\n\nExample:\n```text\n================bm=128, bk=128, bn=128===================\n----- 1024 x 1024 x 1024 -----\nMatmul time:  0.00029766598949208854\nMatmul FLOP/s:  7214407167121.377\nFLOP/s utilization: 3.6621%\n\n----- 4096 x 4096 x 4096 -----\nMatmul time:  0.011771515250438824\nMatmul FLOP/s:  11675553278230.387\nFLOP/s utilization: 5.9267%\n\n----- 8192 x 8192 x 8192 -----\nMatmul time:  0.09183577066054567\nMatmul FLOP/s:  11972585626140.668\nFLOP/s utilization: 6.0775%\n\n================bm=512, bk=1024, bn=1024===================\n----- 1024 x 1024 x 1024 -----\nMatmul time:  0.00012708659982308746\nMatmul FLOP/s:  16897797651282.135\nFLOP/s utilization: 8.5776%\n\n----- 4096 x 4096 x 4096 -----\nMatmul time:  0.00088908776990138\nMatmul FLOP/s:  154584235803001.88\nFLOP/s utilization: 78.4692%\n\n----- 8192 x 8192 x 8192 -----\nMatmul time:  0.006099433819763363\nMatmul FLOP/s:  180264539343531.62\nFLOP/s utilization: 91.5048%\n```\n\nExample:\n```text\nprint(\"================ XLA matmul ===================\")\nmm = jnp.matmul\nanalyze_matmul(1024, 1024, 1024, jnp.bfloat16, mm)\nanalyze_matmul(4096, 4096, 4096, jnp.bfloat16, mm)\nanalyze_matmul(8192, 8192, 8192, jnp.bfloat16, mm)\n```\n\nExample:\n```text\n================ XLA matmul ===================\n----- 1024 x 1024 x 1024 -----\nMatmul time:  0.00011943008983507753\nMatmul FLOP/s:  17981093801113.996\nFLOP/s utilization: 9.1275%\n\n----- 4096 x 4096 x 4096 -----\nMatmul time:  0.0008272899803705514\nMatmul FLOP/s:  166131533963991.34\nFLOP/s utilization: 84.3307%\n\n----- 8192 x 8192 x 8192 -----\nMatmul time:  0.006047147869830951\nMatmul FLOP/s:  181823175395037.44\nFLOP/s utilization: 92.2960%\n```\n\nExample:\n```text\ndef matmul_kernel(x_ref, y_ref, z_ref, acc_ref, *, nsteps, transpose_rhs):\n  @pl.when(pl.program_id(2) == 0)\n  def _():\n    acc_ref[...] = jnp.zeros_like(acc_ref)\n\n  # dot_general expects a data structure (contraction_dims, batch_dims),\n  # where contraction_dims are the set of dimensions for LHS and RHS that will\n  # be contracted (reduced) in the matmul; batch_dims, on the other hand, are\n  # looped over. The remaining dimensions will be the input and output dimension\n  # of the matmul.\n  if transpose_rhs:\n    dims = ((1,), (1,)), ((), ())\n  else:\n    dims = ((1,), (0,)), ((), ())\n\n  acc_ref[...] += jax.lax.dot_general(\n      x_ref[...], y_ref[...], dims, preferred_element_type=jnp.float32,\n  )\n\n  @pl.when(pl.program_id(2) == nsteps - 1)\n  def _():\n    z_ref[...] = acc_ref[...].astype(z_ref.dtype)\n\n\n@jax.jit(static_argnames=['bm', 'bk', 'bn', 'transpose_rhs'])\ndef matmul(\n    x: jax.Array,\n    y: jax.Array,\n    *,\n    bm: int = 128,\n    bk: int = 128,\n    bn: int = 128,\n    transpose_rhs: bool = False,\n):\n  if transpose_rhs:\n    y = y.swapaxes(0, 1)\n    y_block_spec = pl.BlockSpec((bn, bk), lambda i, j, k: (j, k))\n  else:\n    y_block_spec = pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))\n  m, k = x.shape\n  _, n = y.shape\n  return pl.pallas_call(\n      functools.partial(matmul_kernel, nsteps=k // bk, transpose_rhs=transpose_rhs),\n      grid_spec=pltpu.PrefetchScalarGridSpec(\n        num_scalar_prefetch=0,\n        in_specs=[\n            pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)),\n            y_block_spec,\n        ],\n        out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),\n        scratch_shapes=[pltpu.VMEM((bm, bn), jnp.float32)],\n        grid=(m // bm, n // bn, k // bk),\n      ),\n      out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),\n      compiler_params=pltpu.CompilerParams(\n          dimension_semantics=(\"parallel\", \"parallel\", \"arbitrary\")),\n  )(x, y)\n```\n\nExample:\n```text\ndef analyze_matmul(m: int, k: int, n: int, dtype: np.dtype,\n                   mm_func, transpose_rhs: bool = False):\n  x = jnp.ones((m, k), dtype=dtype)\n  if transpose_rhs:\n    y = jnp.ones((n, k), dtype=dtype)\n    @jax.jit\n    def _wrapper(x, y):\n      y = y.swapaxes(0, 1)\n      return mm_func(x, y, transpose_rhs=True)\n  else:\n    y = jnp.ones((k, n), dtype=dtype)\n    _wrapper = mm_func\n  time = benchmark(_wrapper)(x, y)\n  print(f\"----- {m} x {k} x {n} -----\")\n  print(\"Matmul time: \", time)\n  mm_flops = matmul_flops(m, k, n) / time\n  print(\"Matmul FLOP/s: \", mm_flops)\n  print(f\"FLOP/s utilization: {mm_flops / v5e_flops * 100:.4f}%\")\n  print()\n\nprint(\"================bm=128, bk=128, bn=128===================\")\nmm = functools.partial(matmul, bm=128, bk=128, bn=128)\nanalyze_matmul(1024, 1024, 1024, jnp.bfloat16, mm, transpose_rhs=True)\nanalyze_matmul(4096, 4096, 4096, jnp.bfloat16, mm, transpose_rhs=True)\nanalyze_matmul(8192, 8192, 8192, jnp.bfloat16, mm, transpose_rhs=True)\n\nprint(\"================bm=512, bk=1024, bn=1024===================\")\nmm = functools.partial(matmul, bm=512, bk=1024, bn=1024)\nanalyze_matmul(1024, 1024, 1024, jnp.bfloat16, mm, transpose_rhs=True)\nanalyze_matmul(4096, 4096, 4096, jnp.bfloat16, mm, transpose_rhs=True)\nanalyze_matmul(8192, 8192, 8192, jnp.bfloat16, mm, transpose_rhs=True)\n```\n\nExample:\n```text\n================bm=128, bk=128, bn=128===================\n----- 1024 x 1024 x 1024 -----\nMatmul time:  0.0003029372810851783\nMatmul FLOP/s:  7088872126624.065\nFLOP/s utilization: 3.5984%\n\n----- 4096 x 4096 x 4096 -----\nMatmul time:  0.012017967159627005\nMatmul FLOP/s:  11436123235026.848\nFLOP/s utilization: 5.8051%\n\n----- 8192 x 8192 x 8192 -----\nMatmul time:  0.09500920018996112\nMatmul FLOP/s:  11572685861765.383\nFLOP/s utilization: 5.8745%\n\n================bm=512, bk=1024, bn=1024===================\n----- 1024 x 1024 x 1024 -----\nMatmul time:  0.00012131539988331496\nMatmul FLOP/s:  17701657415839.363\nFLOP/s utilization: 8.9856%\n\n----- 4096 x 4096 x 4096 -----\nMatmul time:  0.0008790623804088682\nMatmul FLOP/s:  156347213275211.03\nFLOP/s utilization: 79.3641%\n\n----- 8192 x 8192 x 8192 -----\nMatmul time:  0.006107717020204291\nMatmul FLOP/s:  180020067095253.78\nFLOP/s utilization: 91.3807%\n```\n\nExample:\n```text\ndef matmul_kernel(\n    x_ref, y_ref, z_ref, acc_ref, *, nsteps, transpose_rhs, activation\n):\n  @pl.when(pl.program_id(2) == 0)\n  def _():\n    acc_ref[...] = jnp.zeros_like(acc_ref)\n\n  if transpose_rhs:\n    dims = ((1,), (1,)), ((), ())\n  else:\n    dims = ((1,), (0,)), ((), ())\n\n  acc_ref[...] += jax.lax.dot_general(\n      x_ref[...],\n      y_ref[...],\n      dims,\n      preferred_element_type=jnp.float32,\n  )\n\n  @pl.when(pl.program_id(2) == nsteps - 1)\n  def _():\n    z_ref[...] = activation(acc_ref[...]).astype(z_ref.dtype)\n\n\n@jax.jit(static_argnames=['bm', 'bk', 'bn', 'activation'])\ndef matmul(\n    x: jax.Array,\n    y: jax.Array,\n    *,\n    bm: int = 128,\n    bk: int = 128,\n    bn: int = 128,\n    transpose_rhs: bool = False,\n    activation: Callable[[jax.Array], jax.Array] = lambda x: x,\n):\n  if transpose_rhs:\n    y = y.swapaxes(0, 1)\n    y_block_spec = pl.BlockSpec((bn, bk), lambda i, j, k: (j, k))\n  else:\n    y_block_spec = pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))\n  m, k = x.shape\n  _, n = y.shape\n  return pl.pallas_call(\n      functools.partial(\n          matmul_kernel,\n          nsteps=k // bk,\n          transpose_rhs=transpose_rhs,\n          activation=activation,\n      ),\n      grid_spec=pltpu.PrefetchScalarGridSpec(\n          num_scalar_prefetch=0,\n          in_specs=[\n              pl.BlockSpec((bm, bk), lambda i, j, k: (i, k)),\n              y_block_spec,\n          ],\n          out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),\n          scratch_shapes=[pltpu.VMEM((bm, bn), jnp.float32)],\n          grid=(m // bm, n // bn, k // bk),\n      ),\n      out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),\n      compiler_params=pltpu.CompilerParams(\n          dimension_semantics=(\"parallel\", \"parallel\", \"arbitrary\")),\n  )(x, y)\n```\n\nExample:\n```text\ndef analyze_matmul(m: int, k: int, n: int, dtype: np.dtype,\n                   mm_func, transpose_rhs: bool = False,\n                   activation = lambda x: x):\n  x = jnp.ones((m, k), dtype=dtype)\n  if transpose_rhs:\n    y = jnp.ones((n, k), dtype=dtype)\n    @jax.jit\n    def _wrapper(x, y):\n      y = y.swapaxes(0, 1)\n      return mm_func(x, y, transpose_rhs=True, activation=activation)\n  else:\n    y = jnp.ones((k, n), dtype=dtype)\n    _wrapper = functools.partial(mm_func, activation=activation)\n  time = benchmark(_wrapper)(x, y)\n  print(f\"----- {m} x {k} x {n} -----\")\n  print(\"Matmul time: \", time)\n  mm_flops = matmul_flops(m, k, n) / time\n  print(\"Matmul FLOP/s: \", mm_flops)\n  print(f\"FLOP/s utilization: {mm_flops / v5e_flops * 100:.4f}%\")\n  print()\n\n\nactivation = jax.nn.relu\nprint(\"================bm=128, bk=128, bn=128===================\")\nmm = functools.partial(matmul, bm=128, bk=128, bn=128)\nanalyze_matmul(1024, 1024, 1024, jnp.bfloat16, mm, activation=activation)\nanalyze_matmul(4096, 4096, 4096, jnp.bfloat16, mm, activation=activation)\nanalyze_matmul(8192, 8192, 8192, jnp.bfloat16, mm, activation=activation)\n\nprint(\"================bm=512, bk=1024, bn=1024===================\")\nmm = functools.partial(matmul, bm=512, bk=1024, bn=1024)\nanalyze_matmul(1024, 1024, 1024, jnp.bfloat16, mm, activation=activation)\nanalyze_matmul(4096, 4096, 4096, jnp.bfloat16, mm, activation=activation)\nanalyze_matmul(8192, 8192, 8192, jnp.bfloat16, mm, activation=activation)\n```\n\nExample:\n```text\n================bm=128, bk=128, bn=128===================\n----- 1024 x 1024 x 1024 -----\nMatmul time:  0.00030103540048003196\nMatmul FLOP/s:  7133658182976.541\nFLOP/s utilization: 3.6211%\n\n----- 4096 x 4096 x 4096 -----\nMatmul time:  0.011807117109419778\nMatmul FLOP/s:  11640348122095.826\nFLOP/s utilization: 5.9088%\n\n----- 8192 x 8192 x 8192 -----\nMatmul time:  0.09181861146935262\nMatmul FLOP/s:  11974823079773.941\nFLOP/s utilization: 6.0786%\n\n================bm=512, bk=1024, bn=1024===================\n----- 1024 x 1024 x 1024 -----\nMatmul time:  0.00012622540001757442\nMatmul FLOP/s:  17013086492108.6\nFLOP/s utilization: 8.6361%\n\n----- 4096 x 4096 x 4096 -----\nMatmul time:  0.000896632740041241\nMatmul FLOP/s:  153283442968721.44\nFLOP/s utilization: 77.8089%\n\n----- 8192 x 8192 x 8192 -----\nMatmul time:  0.006130605939542875\nMatmul FLOP/s:  179347953304919.88\nFLOP/s utilization: 91.0396%\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.730Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":25,"totalLines":578,"estimatedTokens":4046}}52{"id":"doc-pallas_core_specific_programming_jax_documentati-8a568797","source":"documentation","title":"Pallas Core-specific Programming — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/tpu/core_map.html","text":"Example:\n```text\nfrom functools import partial\n\nimport jax\nfrom jax.sharding import NamedSharding\nfrom jax.experimental import pallas as pl\nfrom jax.experimental.pallas import tpu as pltpu\nfrom jax.experimental.pallas import tpu_sc as plsc\nimport jax.numpy as jnp\nimport numpy as np\n\n\nnum_devices = jax.local_device_count()\nassert num_devices > 1, \"Please run this notebook with more than one device.\"\n\ntpu_info = pltpu.get_tpu_info()  # This notebook only runs on TPU.\nprint(f\"Running on {num_devices} TPU {tpu_info.chip_version} devices.\")\n```\n\nExample:\n```text\nRunning on 4 TPU v5p devices.\n```\n\nExample:\n```text\n# Mesh of devices\nmesh = jax.make_mesh((jax.device_count(),), ('device',))\nprint(mesh)\n\n# Mesh of cores, within a JAX device\ntc_mesh = pltpu.create_tensorcore_mesh('core')\nprint(tc_mesh)\n\nnum_devices = mesh.size\nnum_cores = len(tc_mesh.devices)\nprint(f\"There are {num_devices} devices, and {num_cores} cores each.\")\n```\n\nExample:\n```text\nMesh('device': 4, axis_types=(Explicit,))\nTensorCoreMesh(devices=array([TensorCore(id=0), TensorCore(id=1)], dtype=object), axis_names=('core',))\nThere are 4 devices, and 2 cores each.\n```\n\nExample:\n```text\n# This runs on every core\ndef swap_cores_kernel(in_hbm, out_hbm,\n                      in_vmem, scratch_vmem, out_vmem,\n                      sem, send_sem, recv_sem):\n  core_index = jax.lax.axis_index('core')\n  num_cores = jax.lax.axis_size('core')\n  slc_size = in_hbm.shape[-1] // num_cores\n  slc = pl.ds(core_index * slc_size, slc_size)\n\n  # Copy in a core-dependent slice of the input\n  pltpu.async_copy(in_hbm.at[:, slc], in_vmem, sem).wait()\n\n  # A barrier to make sure all cores have entered run_scoped.\n  # You won't need this if not doing inter-core communications.\n  dst_core = (core_index + 1) % num_cores\n  sem0 = pltpu.get_barrier_semaphore()\n  pl.semaphore_signal(sem0, 1, device_id={'core': dst_core})\n  pl.semaphore_wait(sem0, 1)\n\n  # Swap data between core 0 and core 1\n  the_copy = pltpu.make_async_remote_copy(\n      in_vmem, scratch_vmem, send_sem, recv_sem, device_id={'core': dst_core},\n  )\n  the_copy.start()\n  the_copy.wait()\n\n  # Core-local compute\n  out_vmem[...] = scratch_vmem[...] * 2\n\n  # Copy out the output\n  pltpu.async_copy(out_vmem, out_hbm.at[:, slc], sem).wait()\n```\n\nExample:\n```text\ninput_shape = (32, 256)\nlocal_vmem_shape = (32 // num_devices, 256 // num_cores)\nin_spec = jax.P('device', None)\nsharding = NamedSharding(mesh, in_spec)\n\n@jax.jit\n@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec,\n         check_vma=False)\ndef swap_cores(x):\n  # Get buffers out of the input and output\n  x_hbm_ref = jax.new_ref(x)\n  o_hbm_ref = jax.new_ref(jax.lax.empty(x.shape, x.dtype))\n\n  @pl.core_map(tc_mesh, compiler_params=pltpu.CompilerParams(collective_id=0))\n  def _():\n    pl.run_scoped(\n        partial(swap_cores_kernel, x_hbm_ref, o_hbm_ref),\n        *([pltpu.VMEM(local_vmem_shape, x.dtype)] * 3),  # VMEM allocations\n        *([pltpu.SemaphoreType.DMA] * 3),                # semaphores\n    )\n  return o_hbm_ref[...]\n\n\nx = jax.random.normal(jax.random.key(0), input_shape, jnp.float32)\nx = jax.device_put(x, sharding)\ny = swap_cores(x)\n\nnp.testing.assert_array_equal(y[:, 128:], x[:, :128] * 2)\nnp.testing.assert_array_equal(y[:, :128], x[:, 128:] * 2)\n```\n\nExample:\n```text\n@jax.jit\n@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec, check_vma=False)\ndef swap_cores(x):\n  scratch_types = [pltpu.VMEM(local_vmem_shape, x.dtype)] * 3 + [pltpu.SemaphoreType.DMA] * 3\n  return pl.kernel(swap_cores_kernel, out_type=x, mesh=tc_mesh,\n                   scratch_types=scratch_types,\n                   compiler_params=pltpu.CompilerParams(collective_id=0))(x)\n\ny = swap_cores(x)\nnp.testing.assert_array_equal(y[:, 128:], x[:, :128] * 2)\nnp.testing.assert_array_equal(y[:, :128], x[:, 128:] * 2)\n```\n\nExample:\n```text\ndef add_one_body(in_vmem, out_vmem):\n  out_vmem[...] = in_vmem[...] + 1\n\ninput_shape = (1024, 1024)\nin_spec = jax.P('device', None)\n\ndef add_one_kernel(x_hbm_ref, o_hbm_ref):\n  in_shape = x_hbm_ref.shape\n  pltpu.emit_pipeline(\n      add_one_body,\n      grid=(in_shape[0] // 8, in_shape[1] // 128),\n      in_specs=[pl.BlockSpec(\n          block_shape=(8, 128), index_map=lambda i, j: (i, j),\n      )],\n      out_specs=[pl.BlockSpec(\n          block_shape=(8, 128), index_map=lambda i, j: (i, j),\n      )],\n      core_axis_name='core',\n      dimension_semantics=(pltpu.PARALLEL, pltpu.ARBITRARY),\n  )(x_hbm_ref, o_hbm_ref)\n\n\n@jax.jit\n@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec, check_vma=False)\ndef add_one(x):\n  return pl.kernel(add_one_kernel, out_type=x, mesh=tc_mesh, scratch_types=[])(x)\n\n\nx = jax.random.normal(jax.random.key(0), input_shape, jnp.float32)\nx = jax.device_put(x, NamedSharding(mesh, in_spec))\ny = add_one(x)\n\nnp.testing.assert_array_equal(y, x + 1)\n```\n\nExample:\n```text\ninput_shape = (1024, 1024)\nin_spec = jax.P('device', None)\noutput_shape = (1024, 512)\n\ndef indexed_add_one_kernel(in_refs, out_refs, i_smem_ref):\n  (x_hbm_ref, i_hbm_ref), o_hbm_ref = in_refs, out_refs\n  in_shape = x_hbm_ref.shape\n  pltpu.sync_copy(i_hbm_ref, i_smem_ref)\n\n  core_idx = jax.lax.axis_index('core')\n  core_slc_size = in_shape[0] // num_cores\n  i_map = lambda i: core_idx * core_slc_size // 8 + i  # split work among cores\n  j_map = lambda j: i_smem_ref[0] // 128 + j           # use the prefetched offset\n\n  pltpu.emit_pipeline(\n      add_one_body,\n      grid=(core_slc_size // 8, output_shape[1] // 128),\n      in_specs=[pl.BlockSpec(\n          block_shape=(8, 128), index_map=lambda i, j: (i_map(i), j_map(j)),\n      )],\n      out_specs=[pl.BlockSpec(\n          block_shape=(8, 128), index_map=lambda i, j: (i_map(i), j),\n      )]\n  )(x_hbm_ref, o_hbm_ref)\n\n\n@jax.jit\n@partial(jax.shard_map, mesh=mesh,\n         in_specs=(in_spec, jax.P()), out_specs=in_spec, check_vma=False)\ndef indexed_add_one(x, index):\n  out_type = jax.ShapeDtypeStruct((x.shape[0], x.shape[1] // 2), x.dtype)\n  return pl.kernel(indexed_add_one_kernel,\n                   out_type=out_type, mesh=tc_mesh,\n                   scratch_types=[pltpu.SMEM((1,), jnp.int32)])((x, index))\n\n\nxs = jax.random.normal(jax.random.key(0), input_shape, jnp.float32)\nxs = jax.device_put(xs, NamedSharding(mesh, in_spec))\nidx = 256\ny = indexed_add_one(xs, jnp.array([idx]))\n\nnp.testing.assert_array_equal(y, xs[:, idx:(idx+512)] + 1)\n```\n\nExample:\n```text\nsc_info = pltpu.get_tpu_info().sparse_core\nassert sc_info is not None\nprint(sc_info)\n\nsc_mesh = plsc.VectorSubcoreMesh(\n    core_axis_name=\"core\", subcore_axis_name=\"subcore\",\n    num_cores=sc_info.num_cores\n)\nsc_num_cores = sc_info.num_cores\nsc_num_subcores = sc_info.num_subcores\n```\n\nExample:\n```text\nSparseCoreInfo(num_cores=4, num_subcores=16, num_lanes=8)\n```\n\nExample:\n```text\ninput_shape = (4096, 128)\nSC_REG_OP_SHAPE = (4, 16)\n\ndef sc_add_one_body(in_vmem, out_vmem):\n  @pl.loop(0, in_vmem.shape[0], step=SC_REG_OP_SHAPE[0])\n  def _reg_loop_0(c0):\n    @pl.loop(0, in_vmem.shape[1], step=SC_REG_OP_SHAPE[1])\n    def _reg_loop_1(c1):\n      slc = (pl.ds(c0, SC_REG_OP_SHAPE[0]), pl.ds(c1, SC_REG_OP_SHAPE[1]))\n      out_vmem[slc] = in_vmem[slc] + 1\n\n\ndef sc_add_one_kernel(x_hbm_ref, o_hbm_ref):\n  in_shape = x_hbm_ref.shape\n  core_idx = jax.lax.axis_index('core')\n  subcore_idx = jax.lax.axis_index(\"subcore\")\n  cm_idx = core_idx * sc_num_subcores + subcore_idx  # index on the core_map\n  slc_size = in_shape[0] // (sc_num_subcores * sc_num_cores)\n  index_map = lambda i, j: (\n      pl.ds(pl.multiple_of(cm_idx * slc_size + i * 8, 8), 8), j)\n\n  pltpu.emit_pipeline(\n      sc_add_one_body,\n      grid=(slc_size // 8, in_shape[1] // 128),\n      in_specs=[pl.BlockSpec(\n          block_shape=(pl.BoundedSlice(8), 128), index_map=index_map,\n      )],\n      out_specs=[pl.BlockSpec(\n          block_shape=(pl.BoundedSlice(8), 128), index_map=index_map,\n      )]\n  )(x_hbm_ref, o_hbm_ref)\n\n\n@jax.jit\n@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec, check_vma=False)\ndef sc_add_one(x):\n  return pl.kernel(sc_add_one_kernel, out_type=x, mesh=sc_mesh, scratch_types=[])(x)\n\n\nx = jax.random.randint(jax.random.key(0), input_shape, 0, 64, jnp.int32)\nx = jax.device_put(x, NamedSharding(mesh, in_spec))\ny = sc_add_one(x)\n\nnp.testing.assert_array_equal(y, x + 1)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.731Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":12,"totalLines":281,"estimatedTokens":2074}}53{"id":"doc-scalar_prefetch_and_block_sparse_computation_jax-1f1dfc0e","source":"documentation","title":"Scalar Prefetch and Block-Sparse Computation — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/tpu/sparse.html","text":"Example:\n```text\nimport timeit\nimport numpy as np\nimport jax\nfrom jax import numpy as jnp\nfrom jax import lax\nfrom jax.experimental import checkify\nfrom jax.experimental import pallas as pl\nfrom jax.experimental.pallas import tpu as pltpu\n\nassert \"TPU\" in jax.devices()[0].device_kind, \"Please run this notebook with TPU devices.\"\nprint(\"Running on\", jax.devices()[0].device_kind)\n```\n\nExample:\n```text\nRunning on TPU v5 lite\n```\n\nExample:\n```text\nclass PrefetchScalarGridSpec:\n  def __init__(self,\n    num_scalar_prefetch: int,\n    grid: tuple[int, ...],\n    in_specs: PyTree[BlockSpec],\n    out_specs: PyTree[BlockSpec],\n    scratch_shapes: tuple[MemorySpace, ...]):\n      ...\n```\n\nExample:\n```text\ndef index_map(*grid_indices, *prefetch_refs):\n    ...\n```\n\nExample:\n```text\ndef kernel(*prefetch_refs, *input_refs, *output_refs, *scratch_refs):\n    ...\n```\n\nExample:\n```text\nkernel = pl.pallas_call(...)\nresult = kernel(*prefetch_args, *input_args)\n```\n\nExample:\n```text\ndef dynamic_slice_kernel(indices, x_ref, o_ref):\n  del indices\n  o_ref[...] = x_ref[...]\n\n@checkify.checkify\n@jax.jit(static_argnums=(2,))\ndef block_dynamic_slice(x, starts, sizes):\n  grid_spec = pltpu.PrefetchScalarGridSpec(\n      num_scalar_prefetch=1,\n      grid=(1, 1),\n      in_specs=[pl.BlockSpec(\n          sizes,\n          lambda i, j, block_idx: (block_idx[0], block_idx[1]))],\n      out_specs=pl.BlockSpec(sizes, lambda *_: (0, 0)),\n  )\n\n  kernel = pl.pallas_call(\n    dynamic_slice_kernel,\n    grid_spec=grid_spec,\n    out_shape=jax.ShapeDtypeStruct(shape=sizes, dtype=x.dtype),\n  )\n  # Checkify inserts a runtime assert that starts are divisible by block size.\n  checkify.check(starts[0] % sizes[0] == 0, \"Starts must be divisible by size.\")\n  checkify.check(starts[1] % sizes[1] == 0, \"Starts must be divisible by size.\")\n  block_idx = jnp.array([starts[0] // sizes[0], starts[1] // sizes[1]])\n  return kernel(block_idx, x)\n\nshape = (512, 512)\nx = jnp.reshape(jnp.arange(np.prod(shape), dtype=jnp.int32), shape)\nerr, result = block_dynamic_slice(x, starts=(128, 256), sizes=(128, 128))\nerr.throw()\nref = lax.dynamic_slice(x, start_indices=(128, 256), slice_sizes=(128, 128))\ndiff = jnp.max(jnp.abs(result - ref))\nprint(\"Error |result - lax.dynamic_slice| =\", diff)\n```\n\nExample:\n```text\nError |result - lax.dynamic_slice| = 0\n```\n\nExample:\n```text\ndef generate_block_sparse_mat(key, M, N, blk_M, blk_N, p=0.2, dtype=jnp.float32):\n  \"\"\"Returns a sampled matrix and its block-sparse representation.\n\n  Args:\n    key: RNG Key.\n    M: Major array dimension.\n    N: Minor array dimension.\n    blk_M: Block size along M dimension.\n    blk_N: Block size along N dimension.\n    p: Probability that a block will be non-zero.\n    dtype: dtype of the sampled matrix.\n\n  Returns:\n    dense_mat: A (M, N) dense sampled array.\n    block_data: A (num_blocks, blk_M, blk_N) array of data blocks representing\n      the non-zero blocks of the matrix.\n    indices_i: A (num_blocks,) array of block indices for the first axis.\n    indices_j: A (num_blocks,) array of block indices for the second axis.\n  \"\"\"\n  mask_key, blocks_key = jax.random.split(key)\n  num_blocks = (M // blk_M, N // blk_N)\n  # We first sample a block mask, denoting which blocks are nonzero.\n  block_mask = jax.random.bernoulli(mask_key, p=p, shape=num_blocks)\n  num_blocks = jnp.sum(block_mask)\n  indices = jnp.where(block_mask)\n  # For each non-zero block, we sample a block of random values.\n  block_data = jax.random.uniform(blocks_key,\n                                  shape=(num_blocks, blk_M, blk_N),\n                                  dtype=dtype)\n  # For checking purposes, create the dense version of the sparse matrix.\n  dense_mat = jnp.zeros((M, N), dtype=dtype)\n  for blk in range(num_blocks):\n    idx_i = indices[0][blk]\n    idx_j = indices[1][blk]\n    slice_i = slice(idx_i * blk_M, (idx_i + 1) * blk_M)\n    slice_j = slice(idx_j * blk_N, (idx_j + 1) * blk_N)\n    dense_mat = dense_mat.at[slice_i, slice_j].set(block_data[blk])\n  return dense_mat, block_data, indices[0], indices[1]\n```\n\nExample:\n```text\nM = N = K = 16384\nblk_M = blk_N = blk_K = 512\n\n\ndef dsd_kernel(idxs_i_ref, idxs_k_ref, # Scalar prefetch inputs.\n               x_ref, y_ref, _, o_ref, # Kernel inputs.\n               accum_scratch,\n               ):\n  \"\"\"A DSD (Dense = Sparse @ Dense) matmul kernel.\"\"\"\n  del idxs_k_ref\n  blk_idx = pl.program_id(1)\n  is_start = blk_idx == 0\n  changed_blocks = (idxs_i_ref[blk_idx] != idxs_i_ref[jnp.maximum(blk_idx-1, 0)])\n  @pl.when(is_start | changed_blocks)\n  def _():\n    accum_scratch[...] = jnp.zeros_like(accum_scratch)\n  accum_scratch[...] += jnp.dot(x_ref[0, :, :], y_ref[...], preferred_element_type=jnp.float32)\n\n  next_block_change = (idxs_i_ref[blk_idx] != idxs_i_ref[jnp.minimum(blk_idx+1, num_blocks)])\n  is_end = blk_idx == (num_blocks - 1)\n  @pl.when(is_end | next_block_change)\n  def _():\n    o_ref[...] = accum_scratch[...].astype(o_ref.dtype)\n\n\ndef x_map(j, blk_idx, blk_idxs_i, blk_idxs_k):\n  del j, blk_idxs_i, blk_idxs_k\n  return (blk_idx, 0, 0)\ndef y_map(j, blk_idx, blk_idxs_i, blk_idxs_k):\n  del blk_idxs_i\n  return (blk_idxs_k[blk_idx], j)\ndef o_map(j, blk_idx, blk_idxs_i, blk_idxs_k):\n  del blk_idxs_k\n  return (blk_idxs_i[blk_idx], j)\n\n(X_dense, X_blocks, indices_i, indices_k) = generate_block_sparse_mat(\n    jax.random.key(0), M, K, blk_M, blk_K, p=0.1, dtype=jnp.bfloat16)\nnum_blocks = X_blocks.shape[0]\nY = jax.random.uniform(jax.random.key(1), shape=(K, N), dtype=jnp.bfloat16)\nzeros = jnp.zeros((M, N), dtype=jnp.bfloat16)\nout_shape = jax.ShapeDtypeStruct((M, N), dtype=jnp.bfloat16)\n\ngrid_spec = pltpu.PrefetchScalarGridSpec(\n    num_scalar_prefetch=2,\n    # Note that while num_blocks is static here, Pallas does support\n    # dynamic grid sizes.\n    grid=(N // blk_N, num_blocks),\n    in_specs=[pl.BlockSpec((1, blk_M, blk_K), x_map),\n              pl.BlockSpec((blk_K, blk_N), y_map),\n              # Placeholder for a zeros-array used by input_output_aliases.\n              pl.BlockSpec((blk_M, blk_N), o_map),\n              ],\n    out_specs=pl.BlockSpec((blk_M, blk_N), o_map),\n    scratch_shapes=[pltpu.VMEM((blk_M, blk_N), dtype=jnp.float32)]\n)\nkernel = pl.pallas_call(\n  dsd_kernel,\n  grid_spec=grid_spec,\n  out_shape=out_shape,\n  # We use input-output aliases to zero-out o_ref for blocks that we never\n  # visit. By passing in an array of zeros we avoid having o_ref start with\n  # uninitialized values.\n  input_output_aliases={4: 0},  # Map zeros to o_ref.\n)\nargs = (indices_i, indices_k, X_blocks, Y, zeros)\nresult = kernel(*args)\n\nref = X_dense @ Y\ndiff = jnp.abs(ref - result)\nprint('mean |result - ref|:', jnp.mean(diff))\n```\n\nExample:\n```text\nmean |result - ref|: 0\n```\n\nExample:\n```text\n# Benchmark Sparse Pallas kernel vs reference JAX implementation\n\ndef benchmark(f, ntrials: int = 100):\n  def run(*args, **kwargs):\n    # Compile function first\n    jax.block_until_ready(f(*args, **kwargs))\n    # Time function\n    result = timeit.timeit(lambda: jax.block_until_ready(f(*args, **kwargs)),\n                           number=ntrials)\n    time = result / ntrials\n    return time\n  return run\n\n\nn_trials = 100\n\npallas_impl = lambda *args: kernel(*args)\ntime = benchmark(pallas_impl, n_trials)(indices_i, indices_k, X_blocks, Y, zeros)\nprint(\"Sparse Kernel: %.3f ms (avg over %d trials)\" % (time * 1000, n_trials))\n\nref_impl = jax.jit(lambda x, y: x @ y)\ntime = benchmark(ref_impl, n_trials)(X_dense, Y)\nprint(\"Reference: %.3f ms (avg over %d trials)\" % (time * 1000, n_trials))\n```\n\nExample:\n```text\nSparse Kernel: 8.136 ms (avg over 100 trials)\nReference: 46.953 ms (avg over 100 trials)\n```\n\nExample:\n```text\ndef mask_index_map(prefetch_map, i, j, ...):\n  next_nonzero_block = prefetch_map[i, j]\n  return (next_nonzero_block, 0, 0)\n```\n\nExample:\n```text\ndef sparsify_mask(mask: jax.Array,\n                  block_shape: tuple[int, int]):\n  \"\"\"Preprocesses a mask into a sparse representation.\n\n  Args:\n    mask: A boolean array of shape [M, N]\n    block_shape: The size of a single block.\n\n  Returns:\n    block_mask: A block_shape array of booleans indicating whether a block\n      is all-zeros (0) or contains non-zero elements (1).\n    prefetch_mask: A block_shape array of integers indicating the index of the\n      next non-zero block.\n    mask_data: A (num_blocks, block_shape) array containing\n      the data for non-zero blocks of the mask.\n  \"\"\"\n  M, N = mask.shape\n  bm, bn = block_shape\n\n  block_mask = jnp.zeros((M // bm, N // bn), dtype=mask.dtype)\n  mask_types_finder = []\n  mask_data = []\n\n  next_mask_type_idx = 0\n  prefetch_mask = jnp.zeros_like(block_mask)\n  next_i = (M // bm) - 1\n  next_j = (N // bn) - 1\n  prefetch_i = jnp.zeros_like(block_mask)\n  prefetch_j = jnp.zeros_like(block_mask)\n  for i in range(M // bm, -1, -1):\n    for j in range(N // bn, -1, -1):\n      mask_block = mask[i * bm :(i + 1) * bm,\n                        j * bn :(j + 1) * bn]\n      is_nonzero = jnp.any(mask_block)\n      if is_nonzero:\n        try:\n          type_index = mask_types_finder.index(str(mask_block))\n        except ValueError:\n          type_index = len(mask_types_finder)\n          mask_types_finder.append(str(mask_block))\n          mask_data.append(mask_block)\n        next_mask_type_idx = type_index\n        next_i = i\n        next_j = j\n      else:\n        type_index = -1\n      block_mask = block_mask.at[i, j].set(is_nonzero)\n      prefetch_mask = prefetch_mask.at[i, j].set(next_mask_type_idx)\n      prefetch_i = prefetch_i.at[i, j].set(next_i)\n      prefetch_j = prefetch_j.at[i, j].set(next_j)\n  return block_mask, prefetch_mask, prefetch_i, prefetch_j, jnp.stack(mask_data)\n```\n\nExample:\n```text\nM = N = K = 16384\nblk_M = blk_N = 512\nblk_K = 1024\n\ndef sparse_mask_matmul(\n    block_mask_ref, prefetch_mask, prefetch_i, prefetch_j, # Scalar prefetch inputs.\n    x_ref, y_ref, mask_ref, o_ref,  # Kernel inputs.\n    accum_scratch\n    ):\n  del prefetch_mask, prefetch_i, prefetch_j\n  i, j, k = pl.program_id(0), pl.program_id(1), pl.program_id(2)\n  should_compute = block_mask_ref[i, j] != 0\n  @pl.when(k == 0)\n  def _():\n    o_ref[...] = jnp.zeros_like(o_ref)\n    accum_scratch[...] = jnp.zeros_like(accum_scratch[...])\n\n  # We only compute the output for blocks with non-zero masks.\n  # Otherwise we skip the computation entirely.\n  @pl.when(should_compute)\n  def _():\n    result = jnp.dot(x_ref[...], y_ref[...], preferred_element_type=jnp.float32)\n    accum_scratch[...] += result\n    @pl.when(k == pl.num_programs(2) - 1)\n    def _():\n      o_ref[...] = (mask_ref[0, ...] * accum_scratch[...]).astype(o_ref.dtype)\n\nX = jax.random.normal(jax.random.key(0), shape=(M, K), dtype=jnp.bfloat16)\nY = jax.random.normal(jax.random.key(1), shape=(K, N), dtype=jnp.bfloat16)\nmask = jnp.ones((M, N), dtype=jnp.int32)\nmask = jnp.tril(mask)\nblock_mask, prefetch_mask, prefetch_i, prefetch_j, sparse_mask_data = sparsify_mask(mask, (blk_M, blk_N))\n\ndef x_map(i, j, k, block_mask, prefetch_mask, prefetch_i, prefetch_j):\n  del prefetch_mask, prefetch_j\n  # Zero-out the k index if the mask is zero, to avoid constantly fetching\n  # new blocks in the inner loop for blocks we are skipping.\n  k_fetch = (block_mask[i, j] != 0) * k\n  return (prefetch_i[i, j], k_fetch)\n\ndef y_map(i, j, k, block_mask, prefetch_mask, prefetch_i, prefetch_j):\n  del prefetch_mask, prefetch_i\n  k_fetch = (block_mask[i, j] != 0) * k\n  return (k_fetch, prefetch_j[i, j])\n\ndef mask_map(i, j, k, block_mask, prefetch_mask, *_):\n  del k, block_mask\n  return (prefetch_mask[i, j], 0, 0)\n\ndef o_map(i, j, k, *_):\n  del k\n  return (i, j)\n\ngrid_spec = pltpu.PrefetchScalarGridSpec(\n    num_scalar_prefetch=4,\n    grid=(M // blk_M, N // blk_N, K // blk_K),\n    in_specs=[pl.BlockSpec((blk_M, blk_K), x_map),\n              pl.BlockSpec((blk_K, blk_N), y_map),\n              pl.BlockSpec((1, blk_M, blk_N), mask_map)],\n    out_specs=pl.BlockSpec((blk_M, blk_N), o_map),\n    scratch_shapes=[pltpu.VMEM((blk_M, blk_N), dtype=jnp.float32)]\n)\nkernel = pl.pallas_call(\n  sparse_mask_matmul,\n  grid_spec=grid_spec,\n  out_shape=jax.ShapeDtypeStruct((M, N), jnp.bfloat16),\n)\nargs = (block_mask, prefetch_mask, prefetch_i, prefetch_j, X, Y, sparse_mask_data)\nresult = kernel(*args)\n\nref = mask * (X @ Y)\ndiff = jnp.abs(ref - result)\nprint('mean |result - ref|:', jnp.mean(diff))\n```\n\nExample:\n```text\nmean |result - ref|: 1.0252e-05\n```\n\nExample:\n```text\nn_trials = 100\n\npallas_impl = lambda *args: kernel(*args)\ntime = benchmark(pallas_impl, n_trials)(block_mask, prefetch_mask, prefetch_i, prefetch_j, X, Y, sparse_mask_data)\nprint(\"Sparse Kernel: %.3f ms (avg over %d trials)\" % (time * 1000, n_trials))\n\nref_impl = jax.jit(lambda mask, x, y: mask * (x @ y))\ntime = benchmark(ref_impl, n_trials)(mask, X, Y)\nprint(\"Reference: %.3f ms (avg over %d trials)\" % (time * 1000, n_trials))\n```\n\nExample:\n```text\nSparse Kernel: 28.648 ms (avg over 100 trials)\nReference: 49.988 ms (avg over 100 trials)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.733Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":19,"totalLines":411,"estimatedTokens":3224}}54{"id":"doc-sparsecore_kernel_writing_jax_documentation-bb4bcf65","source":"documentation","title":"SparseCore Kernel Writing — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/tpu/sparsecore.html","text":"Example:\n```text\nfrom functools import partial\nimport jax\nfrom jax.experimental import pallas as pl\nfrom jax.experimental.pallas import tpu as pltpu\nfrom jax.experimental.pallas import tpu_sc as plsc\nimport jax.numpy as jnp\nimport numpy as np\n\nassert pltpu.get_tpu_info().sparse_core is not None, \"No SparseCore found\"\n```\n\nExample:\n```text\n# Quick way to query basic SC info\n\nassert (sc_info := pltpu.get_tpu_info().sparse_core)\nprint(f\"SparseCore info for TPU {pltpu.get_tpu_info().chip_version}:\")\nprint(sc_info)\n```\n\nExample:\n```text\nSparseCore info for TPU 7x:\nSparseCoreInfo(num_cores=2, num_subcores=16, num_lanes=16, dma_granule_size_bytes=64)\n```\n\nExample:\n```text\nscalar_mesh = plsc.ScalarSubcoreMesh(axis_name=\"core\",\n                                     num_cores=sc_info.num_cores)\nprint(scalar_mesh)\n\nvector_mesh = plsc.VectorSubcoreMesh(\n    core_axis_name=\"core\", subcore_axis_name=\"subcore\"\n)\nprint(vector_mesh)\n```\n\nExample:\n```text\nScalarSubcoreMesh(axis_name='core', num_cores=2)\nVectorSubcoreMesh(core_axis_name='core', subcore_axis_name='subcore', num_cores=2, num_subcores=16)\n```\n\nExample:\n```text\n@jax.jit\ndef cumsum(x):\n  @pl.kernel(out_type=x, mesh=scalar_mesh,\n             scratch_types=[pltpu.SMEM((x.shape[1],), x.dtype),\n                            pltpu.SemaphoreType.DMA])\n  def kernel(x_ref, o_ref, tmp_ref, sem):\n    idx = jax.lax.axis_index('core')\n    pltpu.async_copy(x_ref.at[idx], tmp_ref, sem).wait()\n\n    @pl.loop(1, x.shape[1])\n    def _(i):\n      tmp_ref[i] += tmp_ref[i - 1]\n\n    pltpu.async_copy(tmp_ref, o_ref.at[idx], sem).wait()\n\n  return kernel(x)\n\nx_shape = (sc_info.num_cores, sc_info.num_lanes)\nx = jax.random.randint(jax.random.key(0), x_shape, 0, 64, jnp.int32)\nnp.testing.assert_array_equal(cumsum(x), jnp.cumsum(x, axis=1))\n```\n\nExample:\n```text\nSC_REG_OP_SHAPE = (1, sc_info.num_lanes)\ndma_block = (8, 128)\n\n@jax.jit\ndef sc_add_one(x):\n  @pl.kernel(out_type=x, mesh=vector_mesh, scratch_types=[])\n  def sc_add_one_kernel(x_hbm_ref, o_hbm_ref):\n    in_shape = x_hbm_ref.shape\n\n    def sc_add_one_body(in_vmem, out_vmem):\n      @pl.loop(0, in_vmem.shape[0], step=SC_REG_OP_SHAPE[0])\n      def _(c0):\n        @pl.loop(0, in_vmem.shape[1], step=SC_REG_OP_SHAPE[1])\n        def _(c1):\n          slc = (pl.ds(c0, SC_REG_OP_SHAPE[0]), pl.ds(c1, SC_REG_OP_SHAPE[1]))\n          out_vmem.at[*slc][...] = in_vmem.at[*slc][...] + 1\n\n    pltpu.emit_pipeline(\n        sc_add_one_body,\n        grid=(in_shape[0] // dma_block[0], in_shape[1] // dma_block[1]),\n        in_specs=[pl.BlockSpec(block_shape=dma_block,\n                               index_map=lambda i, j: (i, j))],\n        out_specs=[pl.BlockSpec(block_shape=dma_block,\n                                index_map=lambda i, j: (i, j))],\n        core_axis_name=('core', 'subcore'),\n        dimension_semantics=(pltpu.PARALLEL, pltpu.PARALLEL),\n    )(x_hbm_ref, o_hbm_ref)\n  return sc_add_one_kernel(x)\n\nx = jax.random.randint(jax.random.key(0), (4096, 128), 0, 64, jnp.int32)\ny = sc_add_one(x)\nnp.testing.assert_array_equal(y, x + 1)\n```\n\nExample:\n```text\n@jax.jit\ndef tc_add_one(x):\n  return x + 1\nnp.testing.assert_array_equal(tc_add_one(x), jnp.add(x, 1))\n\n@jax.jit\ndef two_add_ones(x):\n  return sc_add_one(x), tc_add_one(x)\njax.tree.map(np.testing.assert_array_equal, two_add_ones(x), (x + 1, x + 1));\n```\n\nExample:\n```text\n%timeit sc_add_one(x).block_until_ready()\n%timeit tc_add_one(x).block_until_ready()\n\n%timeit jax.block_until_ready(two_add_ones(x))\n```\n\nExample:\n```text\n120 µs ± 2.46 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n113 µs ± 5.61 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n199 µs ± 2.24 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n```\n\nExample:\n```text\nbatch_size = 4096\nvalue_dim = 128\ngather_window_size = 128\nnum_steps = 1024\nsc_num_cores, sc_num_subcores = sc_info.num_cores, sc_info.num_subcores\nnum_indices = gather_window_size * sc_num_cores * sc_num_subcores * num_steps\nx = jnp.arange(batch_size * value_dim).reshape(batch_size, value_dim)\nindices = jax.random.randint(jax.random.key(0), (num_indices,), 0, batch_size,\n                             jnp.int32)\n\n\n@jax.jit\ndef gather(x, indices):\n  indices = indices.reshape((1, num_indices))\n  @pl.kernel(out_type=jax.ShapeDtypeStruct((num_indices, value_dim), x.dtype),\n             mesh=vector_mesh)\n  def kernel(x_hbm, i_hbm, o_hbm):\n    def body(i_vmem, o_vmem):\n      pltpu.sync_copy(x_hbm.at[i_vmem.at[0]], o_vmem)  # The gather op\n\n    pltpu.emit_pipeline(\n        body,\n        grid=(num_indices // gather_window_size,),\n        in_specs=[pl.BlockSpec((1, gather_window_size),\n                               index_map=lambda i: (0, i))],\n        out_specs=[pl.BlockSpec((gather_window_size, value_dim),\n                                index_map=lambda i: (i, 0))],\n        core_axis_name='subcore',\n        dimension_semantics=(pltpu.PARALLEL,),\n    )(i_hbm, o_hbm)\n\n  return kernel(x, indices)\n\nout = gather(x, indices)\nnp.testing.assert_array_equal(out, jnp.take(x, indices, axis=0))\n```\n\nExample:\n```text\n@jax.jit\ndef scatter(x, indices):\n  indices = indices.reshape((1, num_indices))\n  @pl.kernel(out_type=jax.ShapeDtypeStruct((batch_size, value_dim), x.dtype),\n             mesh=vector_mesh, scratch_types=[])\n  def kernel(x_hbm, i_hbm, o_hbm):\n    def body(x_vmem, i_vmem):\n      pltpu.sync_copy(x_vmem, o_hbm.at[i_vmem.at[0]])  # The scatter op\n\n    pltpu.emit_pipeline(\n        body,\n        grid=(num_indices // gather_window_size,),\n        in_specs=[pl.BlockSpec((gather_window_size, value_dim),\n                               index_map=lambda i: (i, 0)),\n                  pl.BlockSpec((1, gather_window_size,),\n                               index_map=lambda i: (0, i))],\n        out_specs=[],\n        core_axis_name='subcore',\n        dimension_semantics=(pltpu.PARALLEL,),\n    )(x_hbm, i_hbm)\n\n  return kernel(x, indices)\n\ngathered = jnp.take(x, indices, axis=0)\nout = scatter(gathered, indices)\nnp.testing.assert_array_equal(out, x)\n```\n\nExample:\n```text\n%timeit jax.block_until_ready(gather(x, indices))\n\ngather_tc = jax.jit(lambda x, i: jnp.take(x, i, axis=0))\ngather_tc(x, indices).block_until_ready()\n\n%timeit jax.block_until_ready(gather_tc(x, indices))\n```\n\nExample:\n```text\n4.05 ms ± 2.02 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n18.1 ms ± 5.24 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.734Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":14,"totalLines":220,"estimatedTokens":1599}}55{"id":"doc-custom_vjp_and_nondiff_argnums_update_guide_jax_-d4a2a232","source":"documentation","title":"custom_vjp and nondiff_argnums update guide — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/4008-custom-vjp-update.html","text":"Example:\n```text\nfrom functools import partial\nimport jax\n\n@partial(jax.custom_vjp, nondiff_argnums=(0, 1))\ndef clip_gradient(lo, hi, x):\n  return x  # identity function\n\ndef clip_gradient_fwd(lo, hi, x):\n  return x, None  # no residual values to save\n\ndef clip_gradient_bwd(lo, hi, _, g):\n  return (jnp.clip(g, lo, hi),)\n\nclip_gradient.defvjp(clip_gradient_fwd, clip_gradient_bwd)\n```\n\nExample:\n```text\nimport jax\n\n@jax.custom_vjp  # no nondiff_argnums!\ndef clip_gradient(lo, hi, x):\n  return x  # identity function\n\ndef clip_gradient_fwd(lo, hi, x):\n  return x, (lo, hi)  # save lo and hi values as residuals\n\ndef clip_gradient_bwd(res, g):\n  lo, hi = res\n  return (None, None, jnp.clip(g, lo, hi))  # return None for lo and hi\n\nclip_gradient.defvjp(clip_gradient_fwd, clip_gradient_bwd)\n```\n\nExample:\n```text\nfrom functools import partial\nimport jax\n\n@partial(jax.custom_vjp, nondiff_argnums=(0,))\ndef skip_app(f, x):\n  return f(x)\n\ndef skip_app_fwd(f, x):\n  return skip_app(f, x), None\n\ndef skip_app_bwd(f, _, g):\n  return (g,)\n\nskip_app.defvjp(skip_app_fwd, skip_app_bwd)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.734Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":55,"estimatedTokens":274}}56{"id":"doc-manual_parallelism_with_shard_map_jax_documentat-ed51a189","source":"documentation","title":"Manual parallelism with shard_map — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/shard_map.html","text":"Example:\n```text\nimport os\nos.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=8' # Use 8 CPU devices\n```\n\nExample:\n```text\nfrom functools import partial\n\nimport jax\nimport jax.numpy as jnp\n\nfrom jax.sharding import Mesh, PartitionSpec as P\nExplicit = jax.sharding.AxisType.Explicit\nAuto = jax.sharding.AxisType.Auto\n```\n\nExample:\n```text\nmesh = jax.make_mesh((4, 2), ('x', 'y'))\njax.set_mesh(mesh)\n\na = jax.device_put(jnp.arange( 8 * 16.).reshape(8, 16), P('x', 'y'))\nb = jax.device_put(jnp.arange(16 *  4.).reshape(16, 4), P('y', None))\n\n@jax.shard_map(in_specs=(P('x', 'y'), P('y', None)), out_specs=P('x', None))\ndef matmul_basic(a_block, b_block):\n  # a_block: f32[2, 8]\n  # b_block: f32[8, 4]\n  c_partialsum = jnp.dot(a_block, b_block)\n  c_block = jax.lax.psum(c_partialsum, 'y')\n  # c_block: f32[2, 4]\n  return c_block\n\nc = matmul_basic(a, b)   # c: f32[8, 4]\n```\n\nExample:\n```text\nfrom jax.tree_util import tree_map, tree_all\n\ndef allclose(a, b):\n  return tree_all(tree_map(partial(jnp.allclose, atol=1e-2, rtol=1e-2), a, b))\n\nallclose(c, jnp.dot(a, b, out_sharding=P('x', None)))\n```\n\nExample:\n```text\nTrue\n```\n\nExample:\n```text\njax.debug.visualize_array_sharding(c)\n```\n\nExample:\n```text\nCPU 0,1   \n            \n            \n  CPU 2,3   \n            \n            \n  CPU 4,5   \n            \n            \n  CPU 6,7\n```\n\nExample:\n```text\nfrom jax.sharding import NamedSharding\n\na = jax.device_put(a, P('x', 'y'))\nb = jax.device_put(b, P('y', None))\n\n@jax.jit\ndef matmul_reference(a, b):\n  return jnp.dot(a, b, out_sharding=P('x', None))\n\nc_ref = matmul_reference(a, b)\nallclose(c_ref, jnp.dot(a, b, out_sharding=P('x', None)))\n```\n\nExample:\n```text\nprint('a blocks:'); jax.debug.visualize_array_sharding(a)\nprint('b blocks:'); jax.debug.visualize_array_sharding(b)\nprint('c blocks:'); jax.debug.visualize_array_sharding(c)\n```\n\nExample:\n```text\na blocks:\nb blocks:\nc blocks:\n```\n\nExample:\n```text\nCPU 0                    CPU 1          \n                                                  \n                                                  \n          CPU 2                    CPU 3          \n                                                  \n                                                  \n          CPU 4                    CPU 5          \n                                                  \n                                                  \n          CPU 6                    CPU 7\n```\n\nExample:\n```text\nCPU 0,2,4,6\n           \n           \n           \n           \n           \nCPU 1,3,5,7\n```\n\nExample:\n```text\ndef check_vmap(f, xs):\n  ans = jax.vmap(f, in_axes=(0,), out_axes=0)(xs)\n  expected = jnp.stack([f(x) for x in xs])  # vmap reference semantics\n  print(allclose(ans, expected))\n\ncheck_vmap(lambda x: x @ x, jnp.arange(12).reshape(4, 3))\n```\n\nExample:\n```text\nimport numpy as np\ndevices = np.array(jax.devices()[:4])\nmesh = Mesh(devices, ('i',))  # mesh.shape['i'] = 4\njax.set_mesh(mesh)\n\ndef check_shmap(f, y):\n  ans = jax.shard_map(f, in_specs=P('i'), out_specs=P('i'))(y)\n  expected = jnp.concatenate([f(y_blk) for y_blk in jnp.split(y, mesh.shape['i'])])\n  print(allclose(ans, expected))\n\ncheck_shmap(lambda x: x.T @ x, jnp.arange(32).reshape(8, 4))\n```\n\nExample:\n```text\nmesh = jax.make_mesh((4, 2), ('i', 'j'))\njax.set_mesh(mesh)\n\n@jax.shard_map(in_specs=P('i', None), out_specs=P('i', 'j'))\ndef f1(x_block):\n  print(x_block.shape)  # prints (3, 12)\n  return x_block\n\nx1 = jax.device_put(jnp.arange(12 * 12).reshape(12, 12), P('i', None))\ny = f1(x1)\n```\n\nExample:\n```text\n(3, 12)\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=P('i', 'j'), out_specs=P('i', 'j'))\ndef f2(x_block):\n  print(x_block.shape)\n  return x_block\n\nx = jnp.arange(12 * 12).reshape(12, 12)\nx_ = jnp.tile(x, (1, mesh.shape['j']))  # x_ has shape (12, 24)\nx_ = jax.device_put(x, P('i', 'j'))\ny = f2(x_)  # prints (3,12), and f1(x) == f2(x_)\n```\n\nExample:\n```text\n(3, 6)\n```\n\nExample:\n```text\nauto_mesh = jax.make_mesh((4, 2), ('i', 'j'), (Auto, Auto))\nwith jax.set_mesh(auto_mesh):\n  x = jnp.array([[3.]])\n\n  z = jax.shard_map(lambda: x, in_specs=(), out_specs=P('i', 'j'))()\n  print(z)  # prints the same as jnp.tile(x, (4, 2))\n\n  z = jax.shard_map(lambda: x, in_specs=(), out_specs=P('i', None))()\n  print(z)  # prints the same as jnp.tile(x, (4, 1)), or just jnp.tile(x, (4,))\n\n  z = jax.shard_map(lambda: x, in_specs=(), out_specs=P(None, None))()\n  print(z)  # prints the same as jnp.tile(x, (1, 1)), or just x\n```\n\nExample:\n```text\n[[3. 3.]\n [3. 3.]\n [3. 3.]\n [3. 3.]]\n[[3.]\n [3.]\n [3.]\n [3.]]\n[[3.]]\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=P('i', 'j'), out_specs=P('i', None))\ndef f3(x_block):\n  return jax.lax.psum(x_block, 'j')\n\nx = jax.device_put(jnp.arange(12 * 12).reshape(12, 12), P('i', 'j'))\ny3 = f3(x)\nprint(y3.shape)\n```\n\nExample:\n```text\n(12, 6)\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=P('i', 'j'), out_specs=P(None, 'j'))\ndef f4(x_block):\n  return jax.lax.psum(x_block, 'i')\n\nx = jax.device_put(jnp.arange(12 * 12).reshape(12, 12), P('i', 'j'))\ny4 = f4(x)\nprint(y4.shape)  # (3,12)\n\n\n@jax.shard_map(in_specs=P('i', 'j'), out_specs=P(None, None))\ndef f5(x_block):\n  return jax.lax.psum(x_block, ('i', 'j'))\n\ny5 = f5(x)\nprint(y5.shape)  # (3,6)\n```\n\nExample:\n```text\n(3, 12)\n(3, 6)\n```\n\nExample:\n```text\nmesh = jax.make_mesh((2,), ('i',))\njax.set_mesh(mesh)\n\n@jax.shard_map(in_specs=P('i'), out_specs=P('i'))\ndef f(x):\n  print(x)\n  return 2 * x\n\nx = jax.device_put(jnp.arange(6.), P('i'))\nf(x)\n```\n\nExample:\n```text\nOn cpu:0 at mesh coordinates (i,) = (0,):\n[0. 1. 2.]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[3. 4. 5.]\n```\n\nExample:\n```text\nArray([ 0.,  2.,  4.,  6.,  8., 10.], dtype=float32)\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=P(), out_specs=P())\ndef f(x):\n  print(x)\n  return 2 * x\n\nx = jnp.arange(6.)\nf(x)\n```\n\nExample:\n```text\nOn cpu:0 at mesh coordinates (i,) = (0,):\n[0. 1. 2. 3. 4. 5.]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[0. 1. 2. 3. 4. 5.]\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=P('i'), out_specs=P())\ndef f(x):\n  y = jax.lax.psum(x, 'i')\n  print(y)\n  return y\n\nx = jax.device_put(jnp.arange(6.), P('i'))\nf(x)\n```\n\nExample:\n```text\nOn cpu:0 at mesh coordinates (i,) = (0,):\n[3. 5. 7.]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[3. 5. 7.]\n```\n\nExample:\n```text\nArray([3., 5., 7.], dtype=float32)\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=P('i'), out_specs=P())\ndef f(x):\n  print(jax.typeof(x))  # f32[3]{V:i}\n  y = jax.lax.psum(x, 'i')\n  print(jax.typeof(y))  # f32[3]\n  return y\n\nx = jax.device_put(jnp.arange(6.), P('i'))\nf(x)\n```\n\nExample:\n```text\nfloat32[3]{V:i}\nfloat32[3]\n```\n\nExample:\n```text\nmesh = jax.make_mesh((4, 2), ('i', 'j'))\njax.set_mesh(mesh)\n\n@jax.shard_map(in_specs=P('i', 'j'), out_specs=P('i'))\ndef f(x):\n  print(jax.typeof(x))  # f32[2,2]{V:(i,j)}\n  y = jax.lax.psum(x, 'j')\n  assert jax.typeof(y).manual_axis_type.varying == {'i'}\n  print(jax.typeof(y))  # f32[2,2]{V:i}\n  return y\n\nx = jax.device_put(jnp.arange(8 * 4.).reshape(8, 4), P('i', 'j'))\nf(x)\n```\n\nExample:\n```text\nfloat32[2,2]{V:(i,j)}\nfloat32[2,2]{V:i}\n```\n\nExample:\n```text\nArray([[ 2.,  4.],\n       [10., 12.],\n       [18., 20.],\n       [26., 28.],\n       [34., 36.],\n       [42., 44.],\n       [50., 52.],\n       [58., 60.]], dtype=float32)\n```\n\nExample:\n```text\nmesh = jax.make_mesh((2,), ('i',))\njax.set_mesh(mesh)\n\nx = jax.device_put(jnp.arange(6.), P('i'))\ntry:\n  y = jax.shard_map(lambda x: x, in_specs=P('i'), out_specs=P())(x)\nexcept Exception as e:\n  print(e)\n```\n\nExample:\n```text\nshard_map applied to the function '_rem_singleton' was given out_specs which require replication which can't be statically inferred given the mesh:\n\nThe mesh given has shape (2,) with corresponding axis names ('i',).\n\nout_specs is P() which implies that the corresponding output value is replicated across mesh axis 'i', but could not infer replication over any axes\n\nCheck if these output values are meant to be replicated over those mesh axes. If not, consider revising the corresponding out_specs entries. If so, consider disabling the check by passing the check_vma=False argument to `jax.shard_map`.\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=P(), out_specs=None)\ndef f(x):\n  print(jax.typeof(x))  # f32[6]\n  y = jax.lax.pcast(x, 'i', to='varying')\n  print(jax.typeof(y))  # f32[6]{i}\n\nx = jnp.arange(6.)\nf(x)\n```\n\nExample:\n```text\nfloat32[6]\nfloat32[6]{V:i}\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=(P('i'), P()), out_specs=P('i'))\ndef f(x, y):\n  return x * y\n\nx = jax.device_put(jnp.arange(6.), P('i'))\ny = jnp.arange(3.)\nprint(jax.make_jaxpr(f)(x, y))\n```\n\nExample:\n```text\n{ lambda ; a:f32[6@i] b:f32[3]. let\n    c:f32[6@i] = shard_map[\n      check_vma=True\n      in_specs=(P('i',), P())\n      jaxpr={ lambda ; d:f32[3]{V:i} e:f32[3]. let\n          f:f32[3]{V:i} = pvary[axes=('i',)] e\n          g:f32[3]{V:i} = mul d f\n        in (g,) }\n      mesh=AbstractMesh('i': 2, axis_types=(Explicit,), device_kind=cpu, num_cores=None, platform=cpu)\n      newly_manual_axes=frozenset({'i'})\n      out_specs=(P('i',),)\n    ] a b\n  in (c,) }\n```\n\nExample:\n```text\nmesh = jax.make_mesh((2,), ('i',))\njax.set_mesh(mesh)\n\n@jax.shard_map(in_specs=(P('i'), P()), out_specs=P('i'))\ndef f(x, y):\n  def body(carry, _):\n    c1, c2 = carry\n    return (c2, c1), ()  # swap the carry\n  (x_, y_), _ = jax.lax.scan(body, (x, y), (), length=2)\n  return x_, y_\n\nx = jnp.arange(6.)\ny = jnp.arange(3.)\n\ntry:\n  f(x, y)\nexcept Exception as e:\n  print(e)\n```\n\nExample:\n```text\nin_specs passed to shard_map: P('i',) does not match the specs of the input: P(None,) for arg: float32[6]. `in_specs` is an optional argument so you can omit specifying it and shard_map will infer the in_specs from the arguments. If you want to reshard your inputs, you can use `jax.reshard` on the arguments and then pass those args to shard_map.\n```\n\nExample:\n```text\nmesh = jax.make_mesh((2,), ('i',))\njax.set_mesh(mesh)\n\n@jax.shard_map(in_specs=(P('i'), P()), out_specs=P('i'))\ndef f(x, y):\n  def body(carry, _):\n    c1, c2 = carry\n    return (c2, c1), ()  # swap the carry\n\n  y = jax.lax.pcast(y, 'i', to='varying')  # apply pcast to fix the error\n  (x_, y_), _ = jax.lax.scan(body, (x, y), (), length=2)\n  return x_, y_\n\nx = jax.device_put(jnp.arange(6.), P('i'))\ny = jnp.arange(3.)\n\nf(x, y)\n```\n\nExample:\n```text\n(Array([0., 1., 2., 3., 4., 5.], dtype=float32),\n Array([0., 1., 2., 0., 1., 2.], dtype=float32))\n```\n\nExample:\n```text\nfrom jax.sharding import Mesh\nSpecs = PyTree[PartitionSpec]\n\ndef shard_map(\n    f: Callable, /, *, out_specs: Specs, mesh: Mesh | None = None,\n    in_specs: Specs | None = None,\n    axis_names: collections.abc.Set[AxisName] = set(),\n    check_vma: bool = True,\n) -> Callable:\n  ...\n```\n\nExample:\n```text\nmesh = Mesh(jax.devices(), ('i',))\nx = jnp.arange(16.)\nf_shmapped = jax.shard_map(f, in_specs=P('i'), out_specs=P('i'))\ny = f_shmapped(x)\n```\n\nExample:\n```text\ndef f_shmapped_ref(x):\n  x_blocks = jnp.array_split(x, mesh.shape['i'])\n  y_blocks = [f(x_blk) for x_blk in x_blocks]\n  return jnp.concatenate(y_blocks)\n```\n\nExample:\n```text\ndef f(x_blk):\n  z_blk = f_part1(x_blk)\n  u_blk = collective(z_blk, axis_name)\n  v_blk = f_part2(x_blk, z_blk, u_blk)\n  return v_blk\n```\n\nExample:\n```text\ndef f_shmapped_ref(x):\n  x_blocks = jnp.array_split(x, mesh.shape[0])\n  z_blocks = [f_part1(x_blk) for x_blk in x_blocks]\n  u_blocks = [collective_ref(i, z_blocks) for i in range(len(z_blocks))]\n  v_blocks = [f_part2(x_blk, z_blk, u_blk) for x_blk, z_blk, u_blk\n              in zip(x_blocks, z_blocks, u_blocks)]\n  return jnp.concatenate(v_blocks)\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax import lax\n\nfrom jax.sharding import Mesh, PartitionSpec as P\n```\n\nExample:\n```text\nmesh1d = Mesh(jax.devices()[:4], ('i',))\njax.set_mesh(mesh1d)\n\n@jax.shard_map(mesh=mesh1d, in_specs=P('i'), out_specs=P(None))\ndef f1(x_block):\n  print('BEFORE:\\n', x_block)\n  y_block = jax.lax.psum(x_block, 'i')\n  print('AFTER:\\n', y_block)\n  return y_block\n```\n\nExample:\n```text\nx = jnp.array([3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 1, 2])\ny = f1(x)\nprint('FINAL RESULT:\\n', y)\n```\n\nExample:\n```text\nBEFORE:\n On cpu:0 at mesh coordinates (i,) = (0,):\n[3 1 4 1]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[5 9 2 6]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[5 3 5 8]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[9 7 1 2]\n\nAFTER:\nOn cpu:0 at mesh coordinates (i,) = (0,):\n[22 20 12 17]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[22 20 12 17]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[22 20 12 17]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[22 20 12 17]\nFINAL RESULT:\n [22 20 12 17]\n```\n\nExample:\n```text\ndef psum_ref(_, x_blocks):\n  tot = sum(x_blocks)\n  return [tot] * len(x_blocks)\n```\n\nExample:\n```text\nmesh2d = Mesh(np.array(jax.devices()[:4]).reshape(2, 2), ('i', 'j'))\njax.set_mesh(mesh2d)\n\n@jax.shard_map(mesh=mesh2d, in_specs=P('i', 'j'), out_specs=P(None, 'j'))\ndef f2(x_block):\n  print('BEFORE:\\n', x_block)\n  y_block = jax.lax.psum(x_block, 'i')\n  print('AFTER:\\n', y_block)\n  return y_block\n\ny = f2(jnp.arange(16).reshape(4, 4))\nprint('FINAL RESULT:\\n', y)\n```\n\nExample:\n```text\nBEFORE:\n On cpu:0 at mesh coordinates (i, j,) = (0, 0):\n[[0 1]\n [4 5]]\n\nOn cpu:1 at mesh coordinates (i, j,) = (0, 1):\n[[2 3]\n [6 7]]\n\nOn cpu:2 at mesh coordinates (i, j,) = (1, 0):\n[[ 8  9]\n [12 13]]\n\nOn cpu:3 at mesh coordinates (i, j,) = (1, 1):\n[[10 11]\n [14 15]]\n\nAFTER:\n On cpu:0 at mesh coordinates (i, j,) = (0, 0):\n[[ 8 10]\n [16 18]]\n\nOn cpu:1 at mesh coordinates (i, j,) = (0, 1):\n[[12 14]\n [20 22]]\n\nOn cpu:2 at mesh coordinates (i, j,) = (1, 0):\n[[ 8 10]\n [16 18]]\n\nOn cpu:3 at mesh coordinates (i, j,) = (1, 1):\n[[12 14]\n [20 22]]\n\nFINAL RESULT:\n [[ 8 10 12 14]\n [16 18 20 22]]\n```\n\nExample:\n```text\n@jax.shard_map(mesh=mesh2d, in_specs=P('i', 'j'), out_specs=P(None, None))\ndef f3(x_block):\n  print('BEFORE:\\n', x_block)\n  y_block = jax.lax.psum(x_block, ('i', 'j'))\n  print('AFTER:\\n', y_block)\n  return y_block\n\ny = f3(jnp.arange(16).reshape(4, 4))\nprint('FINAL RESULT:\\n', y)\n```\n\nExample:\n```text\nBEFORE:\n On cpu:0 at mesh coordinates (i, j,) = (0, 0):\n[[0 1]\n [4 5]]\n\nOn cpu:1 at mesh coordinates (i, j,) = (0, 1):\n[[2 3]\n [6 7]]\n\nOn cpu:2 at mesh coordinates (i, j,) = (1, 0):\n[[ 8  9]\n [12 13]]\n\nOn cpu:3 at mesh coordinates (i, j,) = (1, 1):\n[[10 11]\n [14 15]]\n\nAFTER:\n On cpu:0 at mesh coordinates (i, j,) = (0, 0):\n[[20 24]\n [36 40]]\n\nOn cpu:1 at mesh coordinates (i, j,) = (0, 1):\n[[20 24]\n [36 40]]\n\nOn cpu:2 at mesh coordinates (i, j,) = (1, 0):\n[[20 24]\n [36 40]]\n\nOn cpu:3 at mesh coordinates (i, j,) = (1, 1):\n[[20 24]\n [36 40]]\n\nFINAL RESULT:\n [[20 24]\n [36 40]]\n```\n\nExample:\n```text\njax.set_mesh(mesh1d)\n\n@jax.shard_map(mesh=mesh1d, in_specs=P('i'), out_specs=P('i'))\ndef f4(x_block):\n  print('BEFORE:\\n', x_block)\n  y_block = jax.lax.all_gather(x_block, 'i', tiled=True)\n  print('AFTER:\\n', y_block)\n  return y_block\n\nx = jnp.array([3, 9, 5, 2])\ny = f4(x)\nprint('FINAL RESULT:\\n', y)\n```\n\nExample:\n```text\nBEFORE:\n On cpu:0 at mesh coordinates (i,) = (0,):\n[3]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[9]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[5]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[2]\nAFTER:\nOn cpu:0 at mesh coordinates (i,) = (0,):\n[3 9 5 2]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[3 9 5 2]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[3 9 5 2]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[3 9 5 2]\n\nFINAL RESULT:\n [3 9 5 2 3 9 5 2 3 9 5 2 3 9 5 2]\n```\n\nExample:\n```text\n@jax.shard_map(mesh=mesh1d, in_specs=P('i'), out_specs=P('i'))\ndef f5(x_block):\n  print('BEFORE:\\n', x_block)\n  y_block = jax.lax.all_gather(x_block, 'i', tiled=False)\n  print('AFTER:\\n', y_block)\n  return y_block\n\ny = f5(x)\nprint('FINAL RESULT:\\n', y)\n```\n\nExample:\n```text\nBEFORE:\n On cpu:0 at mesh coordinates (i,) = (0,):\n[3]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[9]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[5]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[2]\n\nAFTER:\n On cpu:0 at mesh coordinates (i,) = (0,):\n[[3]\n [9]\n [5]\n [2]]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[[3]\n [9]\n [5]\n [2]]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[[3]\n [9]\n [5]\n [2]]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[[3]\n [9]\n [5]\n [2]]\nFINAL RESULT:\n [[3]\n [9]\n [5]\n [2]\n [3]\n [9]\n [5]\n [2]\n [3]\n [9]\n [5]\n [2]\n [3]\n [9]\n [5]\n [2]]\n```\n\nExample:\n```text\ndef all_gather_ref(_, x_blocks, *, tiled=False):\n  combine = jnp.concatenate if tiled else jnp.stack\n  return [combine(x_blocks)] * len(x_blocks)\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=P('i'), out_specs=P('i'))\ndef f6(x_block):\n  print('BEFORE:\\n', x_block)\n  y_block = jax.lax.psum_scatter(x_block, 'i', tiled=True)\n  print('AFTER:\\n', y_block)\n  return y_block\n\nx = jnp.array([3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 1, 2])\ny = f6(x)\nprint('FINAL RESULT:\\n', y)\n```\n\nExample:\n```text\nBEFORE:\n On cpu:0 at mesh coordinates (i,) = (0,):\n[3 1 4 1]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[5 9 2 6]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[5 3 5 8]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[9 7 1 2]\n\nAFTER:\n On cpu:0 at mesh coordinates (i,) = (0,):\n[22]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[20]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[12]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[17]\n\nFINAL RESULT:\n [22 20 12 17]\n```\n\nExample:\n```text\ndef psum_scatter_ref(i, x_blocks, *, tiled=False):\n  axis_size = len(x_blocks)\n  tot = sum(x_blocks)\n  if tiled:\n    tot = tot.reshape(axis_size, -1, *tot.shape[1:])  # split leading axis\n  return [tot[i] for i in range(tot.shape[0])]\n```\n\nExample:\n```text\ndef psum(x, axis_name):\n  summed_chunk = jax.lax.psum_scatter(x, axis_name)\n  return jax.lax.all_gather(summed_chunk, axis_name)\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=P('i'), out_specs=P('i'))\ndef f7(x_block):\n  sz = jax.lax.axis_size('i')\n  print('BEFORE:\\n', x_block)\n  y_block = jax.lax.ppermute(x_block, 'i', [(i, (i + 1) % sz) for i in range(sz)])\n  print('AFTER:\\n', y_block)\n  return y_block\n\ny = f7(jnp.arange(8))\nprint('FINAL RESULT:\\n', y)\n```\n\nExample:\n```text\nBEFORE:\n On cpu:0 at mesh coordinates (i,) = (0,):\n[0 1]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[2 3]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[4 5]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[6 7]\n\nAFTER:\n On cpu:0 at mesh coordinates (i,) = (0,):\n[6 7]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[0 1]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[2 3]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[4 5]\n\nFINAL RESULT:\n [6 7 0 1 2 3 4 5]\n```\n\nExample:\n```text\ndef ppermute_ref(i, x_blocks, perm):\n  results = [jnp.zeros_like(x_blocks[0])] * len(x_blocks)\n  for src, dst in perm:\n    results[dst] = x_blocks[src]\n  return results\n```\n\nExample:\n```text\ndef psum_scatter(x, axis_name, *, tiled=False):\n  size = jax.lax.axis_size(axis_name)\n  idx = jax.lax.axis_index(axis_name)  # function instance index along axis_name\n  if tiled:\n    x = x.reshape(size, -1, *x.shape[1:])  # split leading axis\n  shift = partial(jax.lax.ppermute, axis_name=axis_name,\n                  perm=[(i, (i - 1) % size) for i in range(size)])\n  for i in range(1, size):\n    update = shift(x[(idx + i) % size])\n    x = x.at[(idx + i + 1) % size].add(update)\n  return x[idx]\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=P('i'), out_specs=P('i'))\ndef f8(x_block):\n  print('BEFORE:\\n', x_block)\n  y_block = psum_scatter(x_block, 'i', tiled=True)\n  print('AFTER:\\n', y_block)\n  return y_block\n\nx = jnp.array([3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 1, 2])\ny = f8(x)\nprint('FINAL RESULT:\\n', y)\n```\n\nExample:\n```text\n@jax.shard_map(mesh=mesh1d, in_specs=P('i'), out_specs=P('i'))\ndef f9(x_block):\n  print('BEFORE:\\n', x_block)\n  y_block = jax.lax.all_to_all(x_block, 'i', split_axis=0, concat_axis=0,\n                               tiled=True)\n  print('AFTER:\\n', y_block)\n  return y_block\n\nx = jnp.array([3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 1, 2])\ny = f9(x)\nprint('FINAL RESULT:\\n', y)\n```\n\nExample:\n```text\nBEFORE:\n On cpu:0 at mesh coordinates (i,) = (0,):\n[3 1 4 1]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[5 9 2 6]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[5 3 5 8]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[9 7 1 2]\n\nAFTER:\n On cpu:0 at mesh coordinates (i,) = (0,):\n[3 5 5 9]\n\nOn cpu:1 at mesh coordinates (i,) = (1,):\n[1 9 3 7]\n\nOn cpu:2 at mesh coordinates (i,) = (2,):\n[4 2 5 1]\n\nOn cpu:3 at mesh coordinates (i,) = (3,):\n[1 6 8 2]\n\nFINAL RESULT:\n [3 5 5 9 1 9 3 7 4 2 5 1 1 6 8 2]\n```\n\nExample:\n```text\ndef all_to_all_ref(_, x_blocks, *, tiled=False):\n  axis_size = len(x_blocks)\n  if tiled:\n    splits = [jnp.array_split(x, axis_size) for x in x_blocks]\n    return [jnp.concatenate(s) for s in zip(*splits)]\n  else:\n    splits = [list(x) for x in x_blocks]\n    return [jnp.stack(s) for s in zip(*splits)]\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\n\nfrom jax.sharding import Mesh, PartitionSpec as P\n```\n\nExample:\n```text\nmesh = Mesh(jax.devices()[:4], ('i',))\njax.set_mesh(mesh)\n\ndef device_put(x, pspec):\n  return jax.device_put(x, NamedSharding(mesh, pspec))\n```\n\nExample:\n```text\nlhs_spec = P('i', None)\nlhs = device_put(jax.random.normal(jax.random.key(0), (8, 8)), lhs_spec)\n```\n\nExample:\n```text\nrhs_spec = P('i', None)\nrhs = device_put(jax.random.normal(jax.random.key(1), (8, 4)), rhs_spec)\n```\n\nExample:\n```text\n@jax.jit\n@jax.shard_map(in_specs=(lhs_spec, rhs_spec),\n               out_specs=rhs_spec)\ndef matmul_allgather(lhs_block, rhs_block):\n  rhs = jax.lax.all_gather(rhs_block, 'i', tiled=True)\n  return lhs_block @ rhs\n```\n\nExample:\n```text\nout = matmul_allgather(lhs, rhs)\nprint(jnp.allclose(out, lhs @ rhs, atol=1e-3, rtol=1e-3))\n```\n\nExample:\n```text\n@jax.jit\n@jax.shard_map(in_specs=(lhs_spec, rhs_spec),\n               out_specs=rhs_spec)\ndef matmul_allgather_overlapped(lhs_block, rhs_block):\n  size = jax.lax.axis_size('i')\n  idx = jax.lax.axis_index('i')\n  shift = partial(jax.lax.ppermute, axis_name='i',\n                  perm=[(i, (i + 1) % size) for i in range(size)])\n\n  B = lhs_block.shape[1] // size\n  lhs_blocks = lambda i: lax.dynamic_slice_in_dim(lhs_block, i * B, B, 1)\n\n  out_block = lhs_blocks(idx) @ rhs_block\n  for i in range(1, size):\n    rhs_block = shift(rhs_block)\n    out_block += lhs_blocks((idx - i) % size) @ rhs_block\n  return out_block\n```\n\nExample:\n```text\nout = matmul_allgather_overlapped(lhs, rhs)\nprint(jnp.allclose(out, lhs @ rhs, atol=1e-3, rtol=1e-3))\n```\n\nExample:\n```text\n@jax.jit\n@jax.shard_map(in_specs=(lhs_spec, rhs_spec),\n               out_specs=rhs_spec)\ndef matmul_allgather_overlapped_bidi(lhs_block, rhs_block):\n  size = jax.lax.axis_size('i')\n  idx = jax.lax.axis_index('i')\n  shift_up = partial(jax.lax.ppermute, axis_name='i',\n                     perm=[(i, (i + 1) % size) for i in range(size)])\n  shift_dn = partial(jax.lax.ppermute, axis_name='i',\n                     perm=[(i, (i - 1) % size) for i in range(size)])\n\n  B = lhs_block.shape[1] // size // 2  # half-size blocks\n  lhs_blocks = lambda i, hi: lax.dynamic_slice_in_dim(lhs_block, (2*i+hi) * B, B, 1)\n\n  def block_matmul(rhs_lo, rhs_hi, i_lo, i_hi):\n    lhs_lo = jnp.pad(lhs_blocks(i_lo, 0), [(0, 0), (0, B)])\n    lhs_hi = jnp.pad(lhs_blocks(i_hi, 1), [(0, 0), (B, 0)])\n    rhs_lo = jnp.pad(rhs_lo, [(0, B), (0, 0)])\n    rhs_hi = jnp.pad(rhs_hi, [(B, 0), (0, 0)])\n    return (lhs_lo + lhs_hi) @ (rhs_lo + rhs_hi)\n\n  rhs_block_lo, rhs_block_hi = jnp.split(rhs_block, 2, axis=0)\n  out_block = block_matmul(rhs_block_lo, rhs_block_hi, idx, idx)\n  for i in range(1, size):\n    rhs_block_lo = shift_up(rhs_block_lo)\n    rhs_block_hi = shift_dn(rhs_block_hi)\n    out_block += block_matmul(rhs_block_lo, rhs_block_hi,\n                              (idx - i) % size, (idx + i) % size)\n  return out_block\n```\n\nExample:\n```text\nout = matmul_allgather_overlapped_bidi(lhs, rhs)\nprint(jnp.allclose(out, lhs @ rhs, atol=1e-3, rtol=1e-3))\n```\n\nExample:\n```text\nlhs_spec = P(None, 'i')\nlhs = device_put(lhs, lhs_spec)\n\nrhs_spec = P('i', None)\nrhs = device_put(rhs, rhs_spec)\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=(lhs_spec, rhs_spec),\n               out_specs=rhs_spec)\ndef matmul_psumscatter(lhs_block, rhs_block):\n  out_summand = lhs_block @ rhs_block\n  return jax.lax.psum_scatter(out_summand, 'i', tiled=True)\n\nout = matmul_psumscatter(lhs, rhs)\nprint(jnp.allclose(out, lhs @ rhs, atol=1e-3, rtol=1e-3))\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=(lhs_spec, rhs_spec),\n               out_specs=rhs_spec)\ndef matmul_psumscatter_overlapped(lhs_block, rhs_block):\n  size = jax.lax.axis_size('i')\n  idx = jax.lax.axis_index('i')\n  shift = partial(jax.lax.ppermute, axis_name='i',\n                  perm=[(i, (i - 1) % size) for i in range(size)])\n  lhs_block = lhs_block.reshape(size, -1, lhs_block.shape[1])  # split 1st axis\n\n  out_summand = lhs_block[(idx + 1) % size] @ rhs_block\n  for i in range(1, size):\n    out_summand = shift(out_summand)\n    out_summand += lhs_block[(idx + i + 1) % size] @ rhs_block\n  return out_summand\n```\n\nExample:\n```text\nout = matmul_psumscatter_overlapped(lhs, rhs)\nprint(jnp.allclose(out, lhs @ rhs, atol=1e-3, rtol=1e-3))\n```\n\nExample:\n```text\n@jax.shard_map(in_specs=(lhs_spec, rhs_spec),\n               out_specs=rhs_spec)\ndef matmul_psumscatter_overlapped_bidi(lhs_block, rhs_block):\n  size = jax.lax.axis_size('i')\n  idx = jax.lax.axis_index('i')\n  shift_up = partial(jax.lax.ppermute, axis_name='i',\n                     perm=[(i, (i + 1) % size) for i in range(size)])\n  shift_dn = partial(jax.lax.ppermute, axis_name='i',\n                     perm=[(i, (i - 1) % size) for i in range(size)])\n\n  B = lhs_block.shape[0] // size // 2  # half-size blocks\n  lhs_blocks = lambda i, hi: lax.dynamic_slice_in_dim(lhs_block, (2*i+hi) * B, B, 0)\n\n  out_summand_lo = lhs_blocks((idx - 1) % size, 0) @ rhs_block\n  out_summand_hi = lhs_blocks((idx + 1) % size, 1) @ rhs_block\n  for i in range(1, size):\n    out_summand_lo = shift_up(out_summand_lo)\n    out_summand_hi = shift_dn(out_summand_hi)\n    out_summand_lo += lhs_blocks((idx - i - 1) % size, 0) @ rhs_block\n    out_summand_hi += lhs_blocks((idx + i + 1) % size, 1) @ rhs_block\n  return jnp.concatenate([out_summand_lo, out_summand_hi])\n```\n\nExample:\n```text\nout = matmul_psumscatter_overlapped_bidi(lhs, rhs)\nprint(jnp.allclose(out, lhs @ rhs, atol=1e-3, rtol=1e-3))\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\n\ndef predict(params, inputs):\n  for W, b in params:\n    outputs = jnp.dot(inputs, W) + b\n    inputs = jax.nn.relu(outputs)\n  return outputs\n\ndef loss(params, batch):\n  inputs, targets = batch\n  predictions = predict(params, inputs)\n  return jnp.mean(jnp.sum((predictions - targets) ** 2, axis=-1))\n```\n\nExample:\n```text\ndef init_layer(key, n_in, n_out):\n  k1, k2 = jax.random.split(key)\n  W = jax.random.normal(k1, (n_in, n_out)) / jnp.sqrt(n_in)\n  b = jax.random.normal(k2, (n_out,))\n  return W, b\n\ndef init(key, layer_sizes, batch_size):\n  key, *keys = jax.random.split(key, len(layer_sizes))\n  params = list(map(init_layer, keys, layer_sizes[:-1], layer_sizes[1:]))\n\n  key, *keys = jax.random.split(key, 3)\n  inputs = jax.random.normal(keys[0], (batch_size, layer_sizes[0]))\n  targets = jax.random.normal(keys[1], (batch_size, layer_sizes[-1]))\n\n  return params, (inputs, targets)\n```\n\nExample:\n```text\nlayer_sizes = [784, 128, 128, 128, 128, 128, 8]\nbatch_size = 32\n\nparams, batch = init(jax.random.key(0), layer_sizes, batch_size)\n```\n\nExample:\n```text\nfrom jax.sharding import Mesh, PartitionSpec as P\n\nmesh = jax.make_mesh((8,), ('batch',))\njax.set_mesh(mesh)\n\n# replicate initial params on all devices, shard data batch over devices\nbatch = jax.device_put(batch, NamedSharding(mesh, P('batch')))\nparams = jax.device_put(params, NamedSharding(mesh, P()))\n\n# adapt the loss function to sum the losses across devices\n@jax.shard_map(out_specs=P())\ndef loss_dp(params, local_batch):\n  inputs, targets = local_batch\n  predictions = predict(params, inputs)  # use reference 'predict`\n  local_loss = jnp.mean(jnp.sum((predictions - targets)**2, axis=-1))\n  return jax.lax.pmean(local_loss, 'batch')\n```\n\nExample:\n```text\nprint(jax.jit(loss)(params, batch))\nprint(jax.jit(loss_dp)(params, batch))\n```\n\nExample:\n```text\n11.9203\n11.9203\n```\n\nExample:\n```text\ndef allclose(a, b):\n  return tree_all(tree_map(partial(jnp.allclose, atol=1e-2, rtol=1e-2), a, b))\n\nprint(allclose(jax.jit(jax.grad(loss))(params, batch),\n               jax.jit(jax.grad(loss_dp))(params, batch)))\n```\n\nExample:\n```text\n# shard data batch *and params* over devices\nmesh = jax.make_mesh((4,), ('batch',))\njax.set_mesh(mesh)\nbatch = jax.device_put(batch, P('batch'))\nparams = jax.device_put(params, P('batch'))\n\n# adapt the prediction function to gather weights just before their use,\n# and to re-gather them on the backward pass (rather than saving them)\n@partial(jax.remat, policy=lambda op, *_, **__: str(op) != 'all_gather')\ndef predict_fsdp(params_frag, inputs):\n  for W_frag, b_frag in params_frag:\n    W = jax.lax.all_gather(W_frag, 'batch', tiled=True)\n    b = jax.lax.all_gather(b_frag, 'batch', tiled=True)\n    outputs = jnp.dot(inputs, W) + b\n    inputs = jax.nn.relu(outputs)\n  return outputs\n\n@jax.shard_map(out_specs=P())\ndef loss_fsdp(local_params, local_batch):\n  inputs, targets = local_batch\n  predictions = predict_fsdp(local_params, inputs)\n  local_loss = jnp.mean(jnp.sum((predictions - targets) ** 2, axis=-1))\n  return jax.lax.pmean(local_loss, 'batch')\n```\n\nExample:\n```text\nrepl_params = jax.device_put(params, P())\nrepl_batch = jax.device_put(batch, P())\nprint(jax.jit(loss)(repl_params, repl_batch))\nprint(jax.jit(loss_fsdp)(params, batch))\n\nprint(allclose(jax.jit(jax.grad(loss))(repl_params, repl_batch),\n               jax.jit(jax.grad(loss_fsdp))(params, batch)))\n```\n\nExample:\n```text\n11.920298\n11.920298\nTrue\n```\n\nExample:\n```text\nmesh = jax.make_mesh((8,), ('feats',))\njax.set_mesh(mesh)\n\nbatch = jax.device_put(batch, NamedSharding(mesh, P(None, 'feats')))\nparams = jax.device_put(params, NamedSharding(mesh, P('feats')))\n\ndef predict_tp(params, inputs):\n  for W, b in params:\n    outputs = gemm_tp(inputs, W, b)\n    inputs = jax.nn.relu(outputs)\n  return outputs\n\n@jax.shard_map(in_specs=(P(None, 'feats'), P('feats', None), P('feats')),\n               out_specs=P(None, 'feats'))\ndef gemm_tp(inputs, W, b):\n  block_result = jnp.dot(inputs, W)\n  return jax.lax.psum_scatter(block_result, 'feats',\n                              scatter_dimension=1, tiled=True) + b\n\ndef loss_tp(params, batch):\n  inputs, targets = batch\n  predictions = predict_tp(params, inputs)\n  return jnp.mean(jnp.sum((predictions - targets) ** 2, axis=-1))  # NOTE psum!\n```\n\nExample:\n```text\nmesh = jax.make_mesh((4, 2), ('batch', 'feats'))\njax.set_mesh(mesh)\n\nbatch = jax.device_put(batch, NamedSharding(mesh, P('batch', 'feats')))\nparams = jax.device_put(params, NamedSharding(mesh, P(('feats', 'batch'))))\n\n# mostly same as previous predict_fsdp definition, except we call gemm_tp\n@partial(jax.remat, policy=lambda op, *_, **__: str(op) != 'all_gather')\ndef predict_fsdp_tp(params_frag, inputs):\n  for W_frag, b_frag in params_frag:\n    W = jax.lax.all_gather(W_frag, 'batch', tiled=True)\n    b = jax.lax.all_gather(b_frag, 'batch', tiled=True)\n    block_result = jnp.dot(inputs, W)\n    outputs = jax.lax.psum_scatter(block_result, 'feats',\n                                   scatter_dimension=1, tiled=True) + b\n    inputs = jax.nn.relu(outputs)\n  return outputs\n\n@jax.shard_map(in_specs=(P(('feats', 'batch')), P('batch', 'feats')),\n               out_specs=P())\ndef loss_fsdp_tp(local_params, local_batch):\n  inputs, targets = local_batch\n  predictions = predict_fsdp_tp(local_params, inputs)\n  sq_err = jax.lax.psum(jnp.sum((predictions - targets) ** 2, axis=-1), 'feats')\n  return jax.lax.pmean(jnp.mean(sq_err), 'batch')\n```\n\nExample:\n```text\nrepl_params = jax.device_put(params, P())\nrepl_batch = jax.device_put(batch, P())\nprint(jax.jit(loss)(repl_params, repl_batch))\nprint(jax.jit(loss_fsdp_tp)(params, batch))\n\nprint(allclose(jax.jit(jax.grad(loss))(repl_params, repl_batch),\n               jax.jit(jax.grad(loss_fsdp_tp))(params, batch)))\n```\n\nExample:\n```text\nL = len(params) - 2        # num layers, excluding first and last\nN = batch_size             # batch size\nF = params[0][0].shape[1]  # num features\n\n# choose some pipeline parameters\nS = 2      # number of stages\nB = 8      # size of each microbatch\nassert L % S == 0, \"S (number of stages) must divide L (number of inner layers)\"\n\n# compute some useful quantities\nM, ragged = divmod(N, B)  # M is number of microbatches\nassert not ragged, \"B (size of each microbatch) must divide total batch size\"\nK, ragged = divmod(M, S)  # K is microbatches per stage\nassert not ragged, \"S (number of stages) must divide number of microbatches\"\nprint(f'{S} stages, {L // S} layer(s) per stage, {L} pipelined layers total')\nprint(f'{B} examples per microbatch, {M} microbatches total')\n```\n\nExample:\n```text\n2 stages, 2 layer(s) per stage, 4 pipelined layers total\n8 examples per microbatch, 4 microbatches total\n```\n\nExample:\n```text\nmesh = Mesh(jax.devices()[:S], ('stages',))\n\ndef predict_pp(params, inputs):\n  (W_first, b_first), inner_params, (W_last, b_last) = params\n  inputs = jax.nn.relu(jnp.dot(inputs, W_first) + b_first)\n  inputs = spmd_pipeline(lambda Wb, x: jax.nn.relu(x @ Wb[0] + Wb[1]),\n                        inner_params, inputs)\n  outputs = jnp.dot(inputs, W_last) + b_last\n  return outputs\n\n@jax.shard_map(in_specs=((P(), P('stages'), P()), P('stages')), out_specs=P())\ndef loss_pp(params, batch):\n  inputs, targets = batch\n  predictions = predict_pp(params, inputs.reshape(K, B, -1)).reshape(K * B, -1)\n  local_loss = jnp.mean(jnp.sum((predictions - targets)**2, axis=-1))\n  return jax.lax.pmean(local_loss, 'stages')\n```\n\nExample:\n```text\ndef spmd_pipeline(fn, stage_params, inputs):\n  stage = jax.lax.axis_index('stages')\n  outputs = jnp.zeros_like(inputs) * jnp.nan\n  state = jnp.zeros((L // S, B, F)) * jnp.nan\n  for i in range(M+L-1):\n    state = state.at[0].set(jnp.where(stage == 0, inputs[i % K], state[0]))\n    state = jax.vmap(fn)(stage_params, state)\n    outputs = outputs.at[(i-L+1) % K].set(jnp.where(stage == S-1, state[-1], outputs[(i-L+1) % K]))\n    state, inputs, outputs = shift(i, state, inputs, outputs)\n  outputs = jax.lax.ppermute(outputs, 'stages', [(i, (i+1) % S) for i in range(S)])\n  return outputs\n\ndef shift(i, state, inputs, outputs):\n  sh = lambda x, d: jax.lax.ppermute(x, 'stages', [(i, (i+d) % S) for i in range(S)])\n  state = jnp.roll(state, +1, axis=0).at[0].set(sh(state[-1], +1))\n  if (i % K) == (-1 % K):\n    inputs = sh(inputs, +1)\n  if ((i-L+1) % K) == (-1 % K):\n    outputs = sh(outputs, +1)\n  return state, inputs, outputs\n```\n\nExample:\n```text\nfirst_params, *inner_params, last_params = params\nWs, bs = zip(*inner_params)\nparams_stacked = jnp.stack(Ws), jnp.stack(bs)\nfirst_params = jax.device_put(first_params, NamedSharding(mesh, P()))\nparams_stacked = jax.device_put(params_stacked, NamedSharding(mesh, P('stages')))\nlast_params = jax.device_put(last_params, NamedSharding(mesh, P()))\nparams_ = first_params, params_stacked, last_params\n\nbatch_ = jax.device_put(batch, NamedSharding(mesh, P('stages')))\n```\n\nExample:\n```text\njax.set_mesh(mesh)\nprint(jax.jit(loss_pp)(params_, batch_))\n```\n\nExample:\n```text\n11.920299\n```\n\nExample:\n```text\n_ = jax.jit(jax.grad(loss_pp))(params_, batch_)   # don't crash\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.738Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":115,"totalLines":1552,"estimatedTokens":8806}}57{"id":"doc-investigating_a_regression_jax_documentation-8d4bd3de","source":"documentation","title":"Investigating a regression — JAX documentation","url":"https://docs.jax.dev/en/latest/investigating_a_regression.html","text":"Example:\n```text\nfor m in 7 8 9; do\n    for d in `seq -w 1 30`; do\n      docker run -v $PWD:/dir --gpus=all ghcr.io/nvidia/jax:nightly-2023-0${m}-${d} /bin/bash /dir/test.sh &> OUT-0${m}-${d}\n    done\n  Done\n```\n\nExample:\n```text\npip install jmp pyvista numpy matplotlib Rtree trimesh jmp termcolor orbax\n  git clone https://github.com/Autodesk/XLB\n  cd XLB\n  export PYTHONPATH=.\n  export CUDA_VISIBLE_DEVICES=0 # only 1 GPU is needed\n\n  python3 examples/performance/MLUPS3d.py 256 200\n```\n\nExample:\n```text\nOUT-07-06:MLUPS: 587.9240990200157\nOUT-07-07:MLUPS: 587.8907972116419\nOUT-07-08:MLUPS: 587.3186499464459\nOUT-07-09:MLUPS: 587.3130127722537\nOUT-07-10:MLUPS: 587.8526619429658\nOUT-07-17:MLUPS: 570.1631097290182\nOUT-07-18:MLUPS: 570.2819775617064\nOUT-07-19:MLUPS: 570.1672213357352\nOUT-07-20:MLUPS: 587.437153685251\nOUT-07-21:MLUPS: 587.6702557143142\nOUT-07-25:MLUPS: 577.3063618431178\nOUT-07-26:MLUPS: 577.2362978080912\nOUT-07-27:MLUPS: 577.2101850145785\nOUT-07-28:MLUPS: 577.0716349809895\nOUT-07-29:MLUPS: 577.4223280707176\nOUT-07-30:MLUPS: 577.2255967221336\nOUT-08-01:MLUPS: 577.277685388252\nOUT-08-02:MLUPS: 577.0137874289354\nOUT-08-03:MLUPS: 577.1333281553946\nOUT-08-04:MLUPS: 577.305012020407\nOUT-08-05:MLUPS: 577.2143988866626\nOUT-08-06:MLUPS: 577.2409145495443\nOUT-08-07:MLUPS: 577.2602819927345\nOUT-08-08:MLUPS: 577.2823738293221\nOUT-08-09:MLUPS: 577.3453199728248\nOUT-08-11:MLUPS: 577.3161423260563\nOUT-08-12:MLUPS: 577.1697775786824\nOUT-08-13:MLUPS: 577.3049883393633\nOUT-08-14:MLUPS: 576.9051978525331\nOUT-08-15:MLUPS: 577.5331743016213\nOUT-08-16:MLUPS: 577.5117505070573\nOUT-08-18:MLUPS: 577.5930698237612\nOUT-08-19:MLUPS: 577.3539885757353\nOUT-08-20:MLUPS: 577.4190113959127\nOUT-08-21:MLUPS: 577.300394253605\nOUT-08-22:MLUPS: 577.4263792037783\nOUT-08-23:MLUPS: 577.4087536357031\nOUT-08-24:MLUPS: 577.1094728438082\nOUT-08-25:  File \"/XLB/examples/performance/MLUPS3d.py\", line 5, in <module>\nOUT-08-26:MLUPS: 537.0164618489928\nOUT-08-27:MLUPS: 536.9545448661609\nOUT-08-28:MLUPS: 536.2887650464874\nOUT-08-29:MLUPS: 536.7178471720636\nOUT-08-30:MLUPS: 536.6978912984252\nOUT-09-01:MLUPS: 536.7030899164106\nOUT-09-04:MLUPS: 536.5339818238837\nOUT-09-05:MLUPS: 536.6507808565617\nOUT-09-06:MLUPS: 536.7144494518315\nOUT-09-08:MLUPS: 536.7376612408998\nOUT-09-09:MLUPS: 536.7798324141778\nOUT-09-10:MLUPS: 536.726157440174\nOUT-09-11:MLUPS: 536.7446210750584\nOUT-09-12:MLUPS: 536.6707332269023\nOUT-09-13:MLUPS: 536.6777936517823\nOUT-09-14:MLUPS: 536.7581523280307\nOUT-09-15:MLUPS: 536.6156273667873\nOUT-09-16:MLUPS: 536.7320935035265\nOUT-09-17:MLUPS: 536.7104991444398\nOUT-09-18:MLUPS: 536.7492269469092\nOUT-09-19:MLUPS: 536.6760131792959\nOUT-09-20:MLUPS: 536.7361260076634\n```\n\nExample:\n```text\n# Execute this script inside the container:\n  # docker run -v $PWD:/dir --gpus=all ghcr.io/nvidia/jax:nightly-2023-08-24 /bin/bash\n  cd /opt/xla-source\n  git remote update\n  cd /opt/jax-source\n  git remote update\n  pip install jmp pyvista numpy matplotlib Rtree trimesh jmp termcolor orbax\n  cd /tmp\n  git clone https://github.com/Autodesk/XLB\n  cd XLB\n\n  for d in `seq -w 24 26`; do\n      for h in `seq -w 0 24`; do\n          echo $m $d $h\n          /bin/bash /dir/test2.sh Aug $d 2023 $h:00:00 &> OUT-08-${d}-$h\n      done\n  done\n```\n\nExample:\n```text\necho \"param: $@\"\n  cd /opt/xla-source\n  git checkout `git rev-list -1 --before=\"$*\" origin/main`\n  git show -q\n  cd /opt/jax-source\n  git checkout `git rev-list -1 --before=\"$*\" origin/main`\n  git show -q\n\n  rm /opt/jax-source/dist/jax*.whl\n  build-jax.sh # The script is in the nightly container\n\n  export PYTHONPATH=.\n  export CUDA_VISIBLE_DEVICES=0 # only 1 GPU is needed\n\n  python3 examples/performance/MLUPS3d.py 256 200\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.739Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":5,"totalLines":126,"estimatedTokens":925}}58{"id":"doc-writing_high_performance_matrix_multiplication_k-7999d7ca","source":"documentation","title":"Writing high-performance matrix multiplication kernels for Blackwell — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/gpu/blackwell_matmul.html","text":"Example:\n```text\ncutlass_profiler --dist=gaussian,mean:0,stddev:1,scale:-1 --output=results.csv --accumulator-type=f32 --m=4096 --k=4096 --n=8192 --kernels='*sm100*' --A=f16 --B=f16 --C=void --D=f16\n```\n\nExample:\n```text\n@dataclasses.dataclass(frozen=True)\nclass TuningConfig:\n  tile_m: int\n  tile_n: int\n  tile_k: int\n  max_concurrent_steps: int\n```\n\nExample:\n```text\ndef matmul0(a, b, config: TuningConfig):\n  dtype = a.dtype\n  m, k = a.shape\n  _, n = b.shape\n  tile_m, tile_n, tile_k = config.tile_m, config.tile_n, config.tile_k\n  swizzle = plgpu.find_swizzle(tile_k * jnp.dtype(dtype).itemsize * 8)\n  swizzle_elems = swizzle // jnp.dtype(dtype).itemsize\n  transforms = (\n      plgpu.TilingTransform((8, swizzle_elems)), plgpu.SwizzleTransform(swizzle)\n  )\n  if m % tile_m != 0:\n    raise ValueError(f\"{m=} must be divisible by {tile_m=}\")\n  if n % tile_n != 0:\n    raise ValueError(f\"{n=} must be divisible by {tile_n=}\")\n  if k % tile_k != 0:\n    raise ValueError(f\"{k=} must be divisible by {tile_k=}\")\n  m_iters = m // tile_m\n  n_iters = n // tile_n\n  k_iters = k // tile_k\n  max_concurrent_steps = config.max_concurrent_steps\n```\n\nExample:\n```text\ndef kernel(a_gmem, b_gmem, out_gmem, acc_tmem, acc_smem, consumed_barriers):\n    mi = lax.axis_index(\"m\")\n    ni = lax.axis_index(\"n\")\n    m_slice = pl.ds(mi * tile_m, tile_m)\n    n_slice = pl.ds(ni * tile_n, tile_n)\n\n    def do_mma(idxs, a_smem, b_smem):\n      (ki,) = idxs\n      arrive_barrier_slot = ki % 2\n      wait_barrier_slot = 1 - arrive_barrier_slot\n      plgpu.tcgen05_mma(\n          acc_tmem,\n          a_smem,\n          b_smem,\n          barrier=consumed_barriers.at[arrive_barrier_slot],\n          accumulate=(ki > 0),\n      )\n      plgpu.barrier_wait(consumed_barriers.at[wait_barrier_slot])\n\n    # Make sure the wait succeeds in the first iteration.\n    plgpu.barrier_arrive(consumed_barriers.at[1])\n    block_kwargs = dict(transforms=transforms, delay_release=1)\n    plgpu.emit_pipeline(\n      do_mma,\n      in_specs=[\n          plgpu.BlockSpec((tile_m, tile_k), lambda ki: (mi, ki), **block_kwargs),\n          plgpu.BlockSpec((tile_k, tile_n), lambda ki: (ki, ni), **block_kwargs),\n      ],\n      grid=(k_iters,),\n      max_concurrent_steps=max_concurrent_steps,\n    )(a_gmem, b_gmem)\n```\n\nExample:\n```text\ndef kernel(...):\n    ...  # compute pipeline as above\n    final_barrier = 1 - (k_iters % 2)\n    plgpu.barrier_wait(consumed_barriers.at[final_barrier])\n    acc_smem[...] = plgpu.async_load_tmem(acc_tmem).astype(dtype)\n    plgpu.commit_smem()\n    plgpu.copy_smem_to_gmem(acc_smem, out_gmem.at[m_slice, n_slice])\n    plgpu.wait_smem_to_gmem(0, wait_read_only=True)\n```\n\nExample:\n```text\ndef matmul0(a, b, config):\n  ... # Setup code from the first snippet\n  def kernel(...):\n    ... # The whole kernel body\n\n  f = plgpu.kernel(\n      kernel,\n      out_shape=jax.ShapeDtypeStruct((m, n), dtype),\n      grid=(m_iters, n_iters),\n      grid_names=(\"m\", \"n\"),\n      scratch_shapes=dict(\n        acc_tmem=plgpu.TMEM((tile_m, tile_n), jnp.float32),\n        acc_smem=plgpu.SMEM((tile_m, tile_n), dtype, transforms=transforms),\n        consumed_barriers=plgpu.Barrier(\n          num_arrivals=1, num_barriers=2, orders_tensor_core=True\n        ),\n      )\n  )\n  return f(a, b)\n```\n\nExample:\n```text\ndef matmul1(a, b, config: TuningConfig):\n  ... # Setup code remains unmodified\n\n  def kernel(a_gmem, b_gmem, out_gmem,\n             a_smem, b_smem, acc_tmem, acc_smem,\n             load_barriers, consumed_barriers, mma_done_barrier):\n    m_index = lax.axis_index(\"m\")\n    n_index = lax.axis_index(\"n\")\n    m_slice = pl.ds(m_index * tile_m, tile_m)\n    n_slice = pl.ds(n_index * tile_n, tile_n)\n\n    @pl.core_map(plgpu.WarpMesh(axis_name=\"warp\"))\n    def _per_warp():\n      warp_id = lax.axis_index(\"warp\")\n\n      @pl.when(warp_id == 0)\n      def _memory():\n        def _loop_body(ki, _):\n          slot = lax.rem(ki, max_concurrent_steps)\n          @pl.when(ki >= max_concurrent_steps)\n          def _():  # Make sure the data has been consumed before overwriting.\n            plgpu.barrier_wait(consumed_barriers.at[slot])\n          k_slice = pl.ds(ki * tile_k, tile_k)\n          plgpu.copy_gmem_to_smem(\n              a_gmem.at[m_slice, k_slice], a_smem.at[slot], load_barriers.at[slot]\n          )\n          plgpu.copy_gmem_to_smem(\n              b_gmem.at[k_slice, n_slice], b_smem.at[slot], load_barriers.at[slot]\n          )\n\n        lax.fori_loop(0, k_iters, _loop_body, None)\n\n      @pl.when(warp_id == 1)\n      def _compute():\n        def _loop_body(ki, _):\n          slot = lax.rem(ki, max_concurrent_steps)\n          plgpu.barrier_wait(load_barriers.at[slot])  # Wait for data to arrive.\n          plgpu.tcgen05_mma(\n              acc_tmem,\n              a_smem.at[slot],\n              b_smem.at[slot],\n              consumed_barriers.at[slot],\n              accumulate=(ki > 0),\n          )\n        lax.fori_loop(0, k_iters, _loop_body, None)\n        plgpu.tcgen05_commit_arrive(mma_done_barrier)\n\n    plgpu.barrier_wait(mma_done_barrier)\n    acc_smem[...] = plgpu.async_load_tmem(acc_tmem).astype(dtype)\n    plgpu.commit_smem()\n    plgpu.copy_smem_to_gmem(acc_smem, out_gmem.at[m_slice, n_slice])\n    plgpu.wait_smem_to_gmem(0, wait_read_only=True)\n```\n\nExample:\n```text\ndef matmul1(a, b, config: TuningConfig):\n  ... # Setup code remains unmodified\n\n  def kernel(...):\n    ... # Kernel code above\n\n  f = plgpu.kernel(\n      kernel,\n      ...,  # Other parameters remain unchanged\n      scratch_shapes=dict(\n        a_smem=plgpu.SMEM(\n            (max_concurrent_steps, tile_m, tile_k), dtype, transforms=transforms\n        ),\n        b_smem=plgpu.SMEM(\n            (max_concurrent_steps, tile_k, tile_n), dtype, transforms=transforms\n        ),\n        acc_tmem=plgpu.TMEM((tile_m, tile_n), jnp.float32),\n        acc_smem=plgpu.SMEM((tile_m, tile_n), dtype, transforms=transforms),\n        load_barriers=plgpu.Barrier(\n            num_arrivals=2, num_barriers=max_concurrent_steps\n        ),\n        consumed_barriers=plgpu.Barrier(\n            num_arrivals=1,\n            num_barriers=max_concurrent_steps,\n            orders_tensor_core=True,\n        ),\n        mma_done_barrier=plgpu.Barrier(\n            num_arrivals=1, num_barriers=1, orders_tensor_core=True\n        ),\n      )\n  )\n  return f(a, b)\n```\n\nExample:\n```text\ndef matmul2(a, b, config):\n  ... # Setup and kernel code\n  f = plgpu.kernel(\n      ...\n      scratch_shapes=dict(\n        ...\n        # Previously: plgpu.SMEM((tile_m, tile_n), dtype, transforms=transforms),\n        acc_smem=plgpu.SMEM(\n            (2, tile_m, config.epilogue_tile_n), dtype, transforms=transforms\n        ),\n        ...\n      )\n  )\n```\n\nExample:\n```text\ndef matmul2(a, b, config):\n  ... # Setup code remains unchanged\n\n  def kernel(...):\n    ... # Compute part remains unchanged\n\n    plgpu.barrier_wait(mma_done_barrier)\n    out_gmem_window = out_gmem.at[m_slice, n_slice]\n    for ni in range(tile_n // config.epilogue_tile_n):\n      acc_smem_ni = acc_smem.at[ni % 2]\n      ni_slice = pl.ds(ni * config.epilogue_tile_n, config.epilogue_tile_n)\n      # Make sure that previous copy is done before we overwrite.\n      plgpu.wait_smem_to_gmem(1, wait_read_only=True)\n      acc_smem_ni[...] = plgpu.async_load_tmem(acc_tmem.at[:, ni_slice]).astype(dtype)\n      plgpu.commit_smem()\n      plgpu.copy_smem_to_gmem(acc_smem_ni, out_gmem_window.at[:, ni_slice])\n    plgpu.wait_smem_to_gmem(0, wait_read_only=True)\n```\n\nExample:\n```text\ndef matmul3(a, b, config):\n  ...  # Setup code\n  cluster_tile_m = 2 * tile_m\n  cluster_tile_n = 2 * tile_n\n  m_iters = m // cluster_tile_m\n  n_iters = n // cluster_tile_n\n  ... # Setup code and kernel\n\n  f = plgpu.kernel(\n      ...\n      grid=(m_iters, n_iters),\n      ...\n      cluster=(2,),\n      cluster_names=(\"cluster\",),\n      scratch_shapes=dict(\n          ...\n          # Previously: plgpu.TMEM((tile_m, tile_n), jnp.float32),\n          acc_tmem=plgpu.TMEM(\n              (tile_m, cluster_tile_n), jnp.float32, collective=True\n          ),\n          ...\n      )\n  )\n```\n\nExample:\n```text\ndef kernel(...):\n    is_lead_block = lax.axis_index(\"cluster\") == 0\n    m_index = lax.axis_index(\"m\")\n    n_index = lax.axis_index(\"n\")\n    m_slice = pl.ds(m_index * cluster_tile_m, cluster_tile_m)\n    n_slice = pl.ds(n_index * cluster_tile_n, cluster_tile_n)\n```\n\nExample:\n```text\n@pl.core_map(plgpu.WarpMesh(axis_name=\"warp\"))\n    def _per_warp():\n      warp_id = lax.axis_index(\"warp\")\n\n      @pl.when(warp_id == 0)\n      def _memory():\n        def _loop_body(ki, _):\n          ...  # Wait for the data to be consumed, as previously.\n          plgpu.copy_gmem_to_smem(\n              ..., collective_axes=\"cluster\", leader_tracked=plgpu.CopyPartition.PARTITIONED(0)\n          )\n          plgpu.copy_gmem_to_smem(\n              ..., collective_axes=\"cluster\", leader_tracked=plgpu.CopyPartition.PARTITIONED(1)\n          )\n        lax.fori_loop(0, k_iters, _loop_body, None)\n\n      @pl.when(jnp.logical_and(warp_id == 1, is_lead_block))\n      def _compute():\n        def _loop_body(ki, _):\n          ...  # Wait for the data to arrive, as previously.\n          plgpu.tcgen05_mma(\n              ...,\n              collective_axis=\"cluster\",\n          )\n        lax.fori_loop(0, k_iters, _loop_body, None)\n        plgpu.tcgen05_commit_arrive(mma_done_barrier, collective_axis=\"cluster\")\n```\n\nExample:\n```text\ndef matmul3(a, b, config):\n  ...\n  def kernel(...):\n    ... # Compute\n\n    plgpu.barrier_wait(mma_done_barrier)\n    out_m_index = m_index * 2 + lax.axis_index(\"cluster\")\n    out_m_slice = pl.ds(out_m_index * tile_m, tile_m)\n    out_gmem_window = out_gmem.at[out_m_slice, n_slice]\n    for ni in range(cluster_tile_n // config.epilogue_tile_n):\n      ...\n\n  ...\n```\n\nExample:\n```text\ndef matmul4(a, b, config):\n  ...\n\n  num_sms = jax.extend.backend.get_default_device().core_count\n  f = plgpu.kernel(\n      ...\n      grid=(num_sms // 2,),\n      grid_names=(\"cluster_grid\",),\n      ...\n  )\n```\n\nExample:\n```text\ndef matmul4(a, b, config):\n  ...\n\n  def kernel(...):\n    is_lead_block = lax.axis_index(\"cluster\") == 0\n\n    @plgpu.nd_loop((m_iters, n_iters), collective_axes=\"cluster_grid\")\n    def _mn_loop(loop_info: plgpu.NDLoopInfo):\n      m_index, n_index = loop_info.index\n      m_slice = ...\n      n_slice = ...\n\n      ...  # Compute + epilogue\n```\n\nExample:\n```text\ndef matmul4(a, b, config):\n  ...\n  def kernel(...):\n    ...\n    def _mn_loop(...):\n      ...\n\n      @pl.core_map(plgpu.WarpMesh(axis_name=\"warp\"))\n      def _per_warp():\n        warp_id = lax.axis_index(\"warp\")\n\n        @pl.when(warp_id == 0)\n        def _memory():\n          def _loop_body(ki, _):\n            slot = lax.rem(ki, max_concurrent_steps)\n            @pl.when(jnp.logical_or(ki >= max_concurrent_steps, loop_info.local_index > 0))\n            def _():  # Make sure the data has been consumed before overwriting.\n              plgpu.barrier_wait(consumed_barriers.at[slot])\n```\n\nExample:\n```text\ndef matmul4(a, b, config):\n  ...\n  def kernel(...):\n    ...\n    def _mn_loop(...):\n      ...  # Compute + epilogue\n      plgpu.wait_load_tmem()  # Load must complete before MMA can overwrite TMEM.\n```\n\nExample:\n```text\ndef matmul5(a, b, config):\n  ...\n\n  f = plgpu.kernel(\n      ...,\n      num_threads=2,\n      thread_name=\"wg\",\n      scratch_shapes=dict(\n          ...\n          # Previously: plgpu.TMEM((tile_m, cluster_tile_n), jnp.float32, collective=True),\n          acc_tmem=plgpu.TMEM(\n              (tile_m, 2 * cluster_tile_n), jnp.float32, collective=True\n          ),\n          ...\n          # mma_done_barrier (now 2 barriers) + a new store_done_barrier (also 2 barriers)\n          # Previously: plgpu.Barrier(num_arrivals=1, num_barriers=1, orders_tensor_core=True),\n          mma_done_barrier=plgpu.Barrier(\n              num_arrivals=1, num_barriers=2, orders_tensor_core=True\n          ),\n          store_done_barrier=plgpu.ClusterBarrier(\n              collective_axes=(\"cluster\",),\n              num_arrivals=1,\n              num_barriers=2,\n              orders_tensor_core=True,\n          ),\n      ),\n  )\n```\n\nExample:\n```text\ndef matmul(a, b, config):\n  ...\n\n  def kernel(a_gmem, b_gmem, out_gmem,\n             a_smem, b_smem, acc_tmem_slots, acc_smem,\n             load_barriers, consumed_barriers, mma_done_barrier, store_done_barrier):\n    wg_idx = lax.axis_index(\"wg\")\n    is_lead_block = ...\n\n    @plgpu.nd_loop(...)\n    def _mn_loop(...):\n      ...\n      acc_slot = lax.rem(loop_info.local_index, jnp.int32(2))\n      acc_tmem = acc_tmem_slots.at[:, pl.ds(acc_slot * cluster_tile_n, cluster_tile_n)]\n\n      ...\n```\n\nExample:\n```text\ndef matmul(a, b, config):\n  ...\n  def kernel(...):\n    ...\n    def _mn_loop(...):\n      acc_slot = ...\n      acc_tmem = ...\n\n      @pl.when(wg_idx == 0)\n      def _compute_wg():\n        @pl.core_map(plgpu.WarpMesh(axis_name=\"warp\"))\n        def _per_warp():\n          warp_id = lax.axis_index(\"warp\")\n\n          @pl.when(warp_id == 0)\n          def _memory():\n            ... # Memory code remains unchanged\n\n          # Wait for store to complete (except for the first two steps).\n          @pl.when(jnp.logical_and(warp_id == 1, loop_info.local_index >= 2))\n          def _wait_store():\n            plgpu.barrier_wait(store_done_barrier.at[acc_slot])\n          @pl.when(jnp.logical_and(warp_id == 1, is_lead_block))\n          def _compute():\n            ... # Compute loop remains unchanged\n            plgpu.tcgen05_commit_arrive(mma_done_barrier.at[acc_slot], collective_axis=\"cluster\")\n```\n\nExample:\n```text\ndef matmul(a, b, config):\n  ...\n  def kernel(...):\n    ...\n    def _mn_loop(...):\n      ... # Compute\n\n      @pl.when(wg_idx == 1)\n      def _store_wg():\n        ... # Unmodified epilogue\n        plgpu.wait_load_tmem()  # Load must complete before we signal.\n        plgpu.barrier_arrive(store_done_barrier.at[acc_slot])\n```\n\nExample:\n```text\ndef matmul(a, b, config):\n  ...\n  def kernel(...):\n    ...\n    # We now only iterate over a 1D loop (but we still split it across clusters).\n    @plgpu.nd_loop((m_iters * n_iters,), collective_axes=\"cluster_grid\")\n    def _mn_loop(loop_info: plgpu.NDLoopInfo):\n      (lin_idx,) = loop_info.index\n      m_index, n_index = plgpu.planar_snake(\n          lin_idx,  # Linear index.\n          (m_iters, n_iters),  # The 2D iteration space.\n          config.grid_minor_dim,  # 0 or 1, indicates the fastest changing dim.\n          config.grid_tile_width,  # The width of tiles along the fastest changing dim.\n      )\n      ... # Rest of the code remains unmodified\n```\n\nExample:\n```text\ndef matmul6(a, b, config: TuningConfig):\n  dtype = a.dtype\n  m, k = a.shape\n  _, n = b.shape\n  tile_m, tile_n, tile_k = config.tile_m, config.tile_n, config.tile_k\n  swizzle = plgpu.find_swizzle(tile_k * jnp.dtype(dtype).itemsize * 8)\n  swizzle_elems = swizzle // jnp.dtype(dtype).itemsize\n  transforms = (\n      plgpu.TilingTransform((8, swizzle_elems)), plgpu.SwizzleTransform(swizzle)\n  )\n  if m % tile_m != 0:\n    raise ValueError(f\"{m=} must be divisible by {tile_m=}\")\n  if n % tile_n != 0:\n    raise ValueError(f\"{n=} must be divisible by {tile_n=}\")\n  if k % tile_k != 0:\n    raise ValueError(f\"{k=} must be divisible by {tile_k=}\")\n  cluster_tile_m = 2 * tile_m\n  cluster_tile_n = 2 * tile_n\n  m_iters = m // cluster_tile_m\n  n_iters = n // cluster_tile_n\n  k_iters = k // tile_k\n  max_concurrent_steps = config.max_concurrent_steps\n\n  def kernel(a_gmem, b_gmem, out_gmem,\n             a_smem, b_smem, acc_tmem, acc_smem,\n             load_barriers, consumed_barriers, mma_done_barrier, store_done_barrier):\n    wg_idx = lax.axis_index(\"wg\")\n    is_lead_block = lax.axis_index(\"cluster\") == 0\n\n    @plgpu.nd_loop((m_iters * n_iters,), collective_axes=\"cluster_grid\")\n    def _mn_loop(loop_info: plgpu.NDLoopInfo):\n      (lin_idx,) = loop_info.index\n      m_index, n_index = plgpu.planar_snake(\n          lin_idx,\n          (m_iters, n_iters),\n          config.grid_minor_dim,\n          config.grid_tile_width,\n      )\n      m_slice = pl.ds(m_index * cluster_tile_m, cluster_tile_m)\n      n_slice = pl.ds(n_index * cluster_tile_n, cluster_tile_n)\n      acc_slot = lax.rem(loop_info.local_index, jnp.int32(2))\n      mn_acc_tmem = acc_tmem.at[:, pl.ds(acc_slot * cluster_tile_n, cluster_tile_n)]\n\n      @pl.when(wg_idx == 0)\n      def _compute_wg():\n        @pl.core_map(plgpu.WarpMesh(axis_name=\"warp\"))\n        def _per_warp():\n          warp_id = lax.axis_index(\"warp\")\n\n          @pl.when(warp_id == 0)\n          def _memory():\n            def _loop_body(ki, _):\n              slot = lax.rem(ki, max_concurrent_steps)\n              @pl.when(jnp.logical_or(ki >= max_concurrent_steps, loop_info.local_index > 0))\n              def _():  # Make sure the data has been consumed before overwriting.\n                plgpu.barrier_wait(consumed_barriers.at[slot])\n              k_slice = pl.ds(ki * tile_k, tile_k)\n              plgpu.copy_gmem_to_smem(\n                  a_gmem.at[m_slice, k_slice], a_smem.at[slot], load_barriers.at[slot],\n                  collective_axes=\"cluster\", leader_tracked=plgpu.CopyPartition.PARTITIONED(0)\n              )\n              plgpu.copy_gmem_to_smem(\n                  b_gmem.at[k_slice, n_slice], b_smem.at[slot], load_barriers.at[slot],\n                  collective_axes=\"cluster\", leader_tracked=plgpu.CopyPartition.PARTITIONED(1)\n              )\n\n            lax.fori_loop(0, k_iters, _loop_body, None)\n\n          # Wait for store to complete (except for the first two steps).\n          @pl.when(jnp.logical_and(warp_id == 1, loop_info.local_index >= 2))\n          def _wait_store():\n            plgpu.barrier_wait(store_done_barrier.at[acc_slot])\n          @pl.when(jnp.logical_and(warp_id == 1, is_lead_block))\n          def _compute():\n            def _loop_body(ki, _):\n              slot = lax.rem(ki, max_concurrent_steps)\n              plgpu.barrier_wait(load_barriers.at[slot])  # Wait for data to arrive.\n              plgpu.tcgen05_mma(\n                  mn_acc_tmem,\n                  a_smem.at[slot],\n                  b_smem.at[slot],\n                  consumed_barriers.at[slot],\n                  accumulate=(ki > 0),\n                  collective_axis=\"cluster\",\n              )\n            lax.fori_loop(0, k_iters, _loop_body, None)\n            plgpu.tcgen05_commit_arrive(\n                mma_done_barrier.at[acc_slot],\n                collective_axis=\"cluster\",\n            )\n\n      @pl.when(wg_idx == 1)\n      def _store_wg():\n        # Ensure that copies from the previous mn step have completed.\n        plgpu.wait_smem_to_gmem(0, wait_read_only=True)\n        plgpu.barrier_wait(mma_done_barrier.at[acc_slot])\n        out_m_index = m_index * 2 + lax.axis_index(\"cluster\")\n        out_m_slice = pl.ds(out_m_index * tile_m, tile_m)\n        out_gmem_window = out_gmem.at[out_m_slice, n_slice]\n        for ni in range(cluster_tile_n // config.epilogue_tile_n):\n          acc_smem_ni = acc_smem.at[ni % 2]\n          ni_slice = pl.ds(ni * config.epilogue_tile_n, config.epilogue_tile_n)\n          # Make sure that previous copy is done before we overwrite.\n          plgpu.wait_smem_to_gmem(1, wait_read_only=True)\n          acc_smem_ni[...] = plgpu.async_load_tmem(mn_acc_tmem.at[:, ni_slice]).astype(dtype)\n          plgpu.commit_smem()\n          plgpu.copy_smem_to_gmem(acc_smem_ni, out_gmem_window.at[:, ni_slice])\n        plgpu.wait_load_tmem()  # Load must complete before we signal.\n        plgpu.barrier_arrive(store_done_barrier.at[acc_slot])\n    plgpu.wait_smem_to_gmem(0, wait_read_only=True)\n\n  num_sms = backend.get_default_device().core_count\n  f = plgpu.kernel(\n      kernel,\n      out_shape=jax.ShapeDtypeStruct((m, n), dtype),\n      grid=(num_sms // 2,),\n      grid_names=(\"cluster_grid\",),\n      cluster=(2,),\n      cluster_names=(\"cluster\",),\n      num_threads=2,\n      thread_name=\"wg\",\n      scratch_shapes=dict(\n          a_smem=plgpu.SMEM(\n              (max_concurrent_steps, tile_m, tile_k), dtype, transforms=transforms\n          ),\n          b_smem=plgpu.SMEM(\n              (max_concurrent_steps, tile_k, tile_n), dtype, transforms=transforms\n          ),\n          acc_tmem=plgpu.TMEM(\n              (tile_m, 2 * cluster_tile_n), jnp.float32, collective=True\n          ),\n          acc_smem=plgpu.SMEM(\n              (2, tile_m, config.epilogue_tile_n), dtype, transforms=transforms\n          ),\n          load_barriers=plgpu.Barrier(\n              num_arrivals=2, num_barriers=max_concurrent_steps\n          ),\n          consumed_barriers=plgpu.Barrier(\n              num_arrivals=1,\n              num_barriers=max_concurrent_steps,\n              orders_tensor_core=True,\n          ),\n          mma_done_barrier=plgpu.Barrier(\n              num_arrivals=1, num_barriers=2, orders_tensor_core=True\n          ),\n          store_done_barrier=plgpu.ClusterBarrier(\n              collective_axes=(\"cluster\",),\n              num_arrivals=1,\n              num_barriers=2,\n              orders_tensor_core=True,\n          ),\n      )\n  )\n  return f(a, b)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.741Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":24,"totalLines":662,"estimatedTokens":5247}}59{"id":"doc-contributing_to_jax_jax_documentation-ff6e2876","source":"documentation","title":"Contributing to JAX — JAX documentation","url":"https://docs.jax.dev/en/latest/contributing.html","text":"Example:\n```text\ngit clone https://github.com/YOUR_USERNAME/jax\ncd jax\npip install -r build/test-requirements.txt  # Installs all testing requirements.\npip install -e \".[cpu]\"  # Installs JAX from the current directory in editable mode.\n```\n\nExample:\n```text\ngit remote add upstream https://github.com/jax-ml/jax.git\n```\n\nExample:\n```text\ngit checkout -b name-of-change\n```\n\nExample:\n```text\npip install pre-commit\npre-commit run --all\n```\n\nExample:\n```text\npytest -n auto tests/\n```\n\nExample:\n```text\nJAX_ENABLE_X64=True pytest -n auto tests/\n```\n\nExample:\n```text\npytest -n auto tests/lax_scipy_test.py\n```\n\nExample:\n```text\npytest -n auto tests/lax_scipy_test.py -k testLogSumExp\n```\n\nExample:\n```text\ngit add file1.py file2.py ...\ngit commit -m \"Your commit message\"\n```\n\nExample:\n```text\ngit fetch upstream\ngit rebase upstream/main\n```\n\nExample:\n```text\ngit push --set-upstream origin name-of-change\n```\n\nExample:\n```text\npip install pre-commit\npre-commit run --all-files\n```\n\nExample:\n```text\nwheel_sources(\n   name = \"jax_sources\",\n   data_srcs = [...],\n   py_srcs = [...],\n   static_srcs = [\n      ...\n      \"//:file.txt\"\n   ],\n)\n```\n\nExample:\n```text\nwheel_sources(\n   name = \"jax_sources\",\n   data_srcs = [\n      ...\n      \"//:cc_target\"\n   ],\n   py_srcs = [...],\n   static_srcs = [...],\n)\n```\n\nExample:\n```text\npytype_strict_library(\n   name = \"init\",\n   srcs = [\"__init__.py\"],\n   visibility = [\"//visibility:public\"],\n)\n```\n\nExample:\n```text\nwheel_sources(\n   name = \"jax_sources\",\n   data_srcs = [...],\n   py_srcs = [\n      ...\n      \"//jax/test_package:init\",\n   ],\n   static_srcs = [...],\n)\n```\n\nExample:\n```text\npytype_strict_library(\n   name = \"example\",\n   srcs = [\"__init__.py\",\n            \"example.py\"],\n   visibility = [\"//visibility:public\"],\n)\n```\n\nExample:\n```text\nwheel_sources(\n   name = \"jax_sources\",\n   data_srcs = [...],\n   py_srcs = [\n      ...\n      \"//jax/test_package:example\",\n   ],\n   static_srcs = [...],\n)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.742Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":18,"totalLines":139,"estimatedTokens":491}}60{"id":"doc-writing_mosaic_gpu_kernels_with_pallas_jax_docum-17e87380","source":"documentation","title":"Writing Mosaic GPU kernels with Pallas — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/gpu/reference.html","text":"Example:\n```text\nimport jax.experimental.pallas as pl\nimport jax.experimental.pallas.mosaic_gpu as plgpu\n```\n\nExample:\n```text\ndef body(..., scratch_ref):\n  # Asynchronous copy will reformat the GMEM data to match the SMEM transforms\n  plgpu.copy_gmem_to_smem(..., scratch_ref, barrier)\n  plgpu.barrier_wait(barrier)\n  plgpu.wgmma(..., scratch_ref)  # wgmma only accepts properly transformed refs\n  ...\n```\n\nExample:\n```text\ntransforms = (plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128))\nf = pl.pallas_call(\n  in_specs=plgpu.BlockSpec(in_block_shape, in_index_map, transforms=transforms),\n  out_specs=plgpu.BlockSpec(out_block_shape, out_index_map, transforms=transforms),\n  ...\n)\n```\n\nExample:\n```text\ntransforms = (plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128))\nf = pl.pallas_call(\n  scratch_shapes=plgpu.SMEM((128, 128), jnp.float16, transforms=transforms),\n  ...\n)\n```\n\nExample:\n```text\ndef mma_transforms(shape_dtype: jax.ShapeDtypeStruct):\n  assert len(shape_dtype.shape) == 2\n  if shape_dtype.shape[0] % 8:\n    raise ValueError(\"Number of rows must be divisible by 8\")\n  for swizzle_bytes in (128, 64, 32):\n    swizzle_elems = swizzle_bytes // shape_dtype.dtype.itemsize\n    if shape_dtype.shape[-1] % swizzle_elems == 0:\n      return (plgpu.TilingTransform((8, swizzle_elems)),\n              plgpu.SwizzleTransform(swizzle_bytes))\n  raise ValueError(\"Failed to find transforms for the specified window type\")\n```\n\nExample:\n```text\nassert acc_ref.shape == (M, N) and a_ref.shape == (K, M) and b_ref.shape == (K, N)\na_ref_t = a_ref.transpose((1, 0))\nassert a_ref_t.shape == (M, K)  # The shape expected by plgpu.wgmma\nplgpu.wgmma(acc, a_ref_t, b_ref)\n```\n\nExample:\n```text\ndef compute(acc_ref):\n  ...\n  return acc_ref[...]\noutput = pl.run_scoped(compute, plgpu.ACC((m, n), jnp.float32))\n```\n\nExample:\n```text\ndef compute(acc_ref):\n  ...\n  return # pl.run_state only returns the final value of the accumulator\noutput = pl.run_state(compute)(plgpu.ACC.init(init_array))\n```\n\nExample:\n```text\nbuffers = 3  # In reality you might want even more\nassert a_smem.shape == (buffers, m, k)\nassert b_smem.shape == (buffers, k, n)\nassert acc_ref.shape == (m, n)\n\ndef fetch_a_b(ki, slot):\n  a_slice = ... # Replace with the right M/K slice\n  b_slice = ... # Replace with the right K/N slice\n  plgpu.copy_gmem_to_smem(a_gmem.at[a_slice], a_smem.at[slot], a_loaded.at[slot])\n  plgpu.copy_gmem_to_smem(b_gmem.at[b_slice], b_smem.at[slot], b_loaded.at[slot])\n\ndef loop_body(i, _):\n  slot = jax.lax.rem(i, buffers)\n  plgpu.barrier_wait(a_loaded.at[slot])\n  plgpu.barrier_wait(b_loaded.at[slot])\n  plgpu.wgmma(acc_ref, a_smem.at[slot], b_smem.at[slot])\n  # We know that only the last issued WGMMA is running, so we can issue a async load in\n  # into the other buffer\n  load_i = i + buffers - 1\n  load_slot = jax.lax.rem(load_i, buffers)\n  @pl.when(jnp.logical_and(load_i >= buffers, load_i < num_steps))\n  def _do_fetch():\n    fetch_a_b(load_i, slot)\nfor slot in range(buffers):\n  fetch_a_b(slot, slot)\njax.lax.fori_loop(0, num_steps, loop_body, None)\n```\n\nExample:\n```text\n@functools.partial(pl.run_scoped, tmem_ref=plgpu.TMEM((128, 128), jnp.float32))\ndef barrier_scope(tmem_ref):\n  ...\n```\n\nExample:\n```text\n@functools.partial(pl.run_scoped,\n                   acc_ref=plgpu.TMEM((128, 128), jnp.float16, packed=False),\n                   lhs_ref=plgpu.TMEM((128, 128), jnp.float16, packed=True))\ndef barrier_scope(acc_ref, lhs_ref):\n  plgpu.tcgen05_mma(acc_ref, lhs_ref, rhs_smem_ref, ...)\n  ...\n```\n\nExample:\n```text\nsmem_ref[...] = plgpu.async_load_tmem(tmem_ref)\nplgpu.commit_smem()\nplgpu.copy_smem_to_gmem(smem_ref, gmem_ref)\nplgpu.wait_smem_to_gmem(0)\nplgpu.wait_load_tmem()  # Wait for the read to fully complete before we overwrite tmem_ref again.\n```\n\nExample:\n```text\nplgpu.async_store_tmem(tmem_ref, smem_ref[...])\nplgpu.commit_tmem()\nsmem_ref2[...] = plgpu.async_load_tmem(tmem_ref)  # Safe to read from tmem_ref now\n```\n\nExample:\n```text\n@functools.partial(pl.run_scoped, barrier_ref=plgpu.Barrier(orders_tensor_core=True))\ndef barrier_scope(barrier_ref):\n  plgpu.tcgen05_mma(acc_tmem, lhs_ref, rhs_ref, barrier_ref, accumulate=False)\n  plgpu.barrier_wait(barrier_ref)\n  # We can read the result now.\n  result = plgpu.async_load_tmem(acc_tmem)\n  ...\n```\n\nExample:\n```text\n@functools.partial(pl.run_scoped, barrier_ref=plgpu.Barrier(orders_tensor_core=True))\ndef barrier_scope(barrier_ref):\n  plgpu.tcgen05_mma(acc_tmem, lhs_ref, rhs_ref, accumulate=False)\n  plgpu.tcgen05_mma(acc_tmem, lhs_ref2, rhs_ref2)\n  plgpu.tcgen05_commit(barrier_ref)\n  plgpu.barrier_wait(barrier_ref)\n  # We can read the result now. Both MMAs have completed.\n  result = plgpu.async_load_tmem(acc_tmem)\n  ...\n```\n\nExample:\n```text\nplgpu.copy_gmem_to_smem(\n    b_gmem,  # [K, N]\n    b_smem,  # [K, N // 2]\n    b_tma_barrier,\n    collective_axes=\"x\",\n    leader_tracked=plgpu.CopyPartition.PARTITIONED(1),\n)\n```\n\nExample:\n```text\n@functools.partial(\n  pl.pallas_call,\n  grid=(2,),\n  in_specs=[pl.BlockSpec(block_shape=(128,), index_map=lambda i: (i,))],\n  out_specs=pl.BlockSpec(block_shape=(128,), index_map=lambda i: (i,)),\n  out_shape=jax.ShapeDtypeStruct((256,), jnp.float32), # Total output shape\n)\ndef run_kernel(x_ref, y_ref):\n  # x_ref and y_ref are in SMEM!\n  y_ref[...] = x_ref[...] + 1\n\nx = jnp.arange(256, dtype=jnp.float32)\ny = run_kernel(x)\nnp.testing.assert_array_equal(y, x + 1)\n```\n\nExample:\n```text\n@pl.run_state\ndef run_kernel(refs):\n  x_ref, y_ref = refs\n  # Here, we're not in the kernel yet! pl.run_state simply changes the JAX\n  # immutable arrays into mutable GMEM (not SMEM!) references.\n\n  # Define the mesh: 2 CUDA blocks over 1 axis called \"x\"\n  mesh = plgpu.Mesh(grid=(2,), grid_names=(\"x\",))\n\n  @pl.core_map(mesh)  # core_map executes the body\n  def kernel_body():\n    # Once we enter the pl.core_map scope, we are in the body of the kernel.\n    block_slice = pl.ds(jax.lax.axis_index(\"x\") * 128, 128)\n    y_ref[block_slice] = x_ref[block_slice] + 1\n\nx = jnp.arange(256, dtype=jnp.float32)\ny_init = jnp.zeros_like(x)\n_, y = run_kernel((x, y_init))\nnp.testing.assert_array_equal(y, x + 1)\n```\n\nExample:\n```text\n@functools.partial(\n    plgpu.kernel,\n    out_shape=jax.ShapeDtypeStruct((256,), jnp.float32),\n    grid=(2,),\n    grid_names=(\"x\",),\n)\ndef run_kernel(x_ref, y_ref):\n  # x_ref and y_ref are in GMEM!\n  block_slice = pl.ds(jax.lax.axis_index(\"x\") * 128, 128)\n  y_ref[block_slice] = x_ref[block_slice] + 1\n\nx = jnp.arange(256, dtype=jnp.float32)\ny = run_kernel(x)  # No need to preallocate outputs as in pl.core_map.\nnp.testing.assert_array_equal(y, x + 1)\n```\n\nExample:\n```text\nx = jnp.arange(128, dtype=jnp.float32)\n\n@functools.partial(\n    plgpu.kernel,\n    out_shape=x,\n    scratch_shapes=dict(\n        smem_ref=plgpu.SMEM(x.shape, x.dtype),\n        barrier_ref=plgpu.Barrier(),\n    ),\n    num_threads=2,\n    thread_name=\"pallas_thread\",\n)\ndef run_kernel(x_ref, y_ref, smem_ref, barrier_ref):\n  thread_id = jax.lax.axis_index(\"pallas_thread\")\n\n  @pl.when(thread_id == 0)\n  def producer_thread():\n    smem_ref[...] = x_ref[...] + 1\n    plgpu.barrier_arrive(barrier_ref)  # Signal the consumer thread\n\n  @pl.when(thread_id == 1)\n  def consumer_thread():\n    plgpu.barrier_wait(barrier_ref)  # Wait for the producer thread\n    out_ref[...] = smem_ref[...] + 1\n\ny = run_kernel(x)  # There's no need to preallocate the input anymore.\nnp.testing.assert_array_equal(y, x + 2)\n```\n\nExample:\n```text\n@functools.partial(\n    plgpu.kernel,\n    out_shape=jax.ShapeDtypeStruct((2, 128), jnp.float32),\n    scratch_shapes=dict(\n        smem_ref=plgpu.SMEM((128,), jnp.float32),\n        barrier_ref=plgpu.Barrier(),\n    ),\n    cluster=(2,),\n    cluster_names=(\"cluster\",),\n)\ndef run_kernel(x_ref, y_ref, smem_ref, barrier_ref):\n  # Specifying collective_axes will enable TMA multicast automatically.\n  plgpu.copy_gmem_to_smem(x_ref, smem_ref, barrier_ref, collective_axes=\"cluster\")\n  plgpu.barrier_wait(barrier_ref)\n  plgpu.copy_smem_to_gmem(smem_ref, o_ref.at[jax.lax.axis_index(\"cluster\")])\n  plgpu.wait_smem_to_gmem(0)\n\nx = jnp.arange(128, dtype=jnp.float32)\ny = run_kernel(x)\n# Each block gets the same data and writes it out.\nnp.testing.assert_array_equal(y, jnp.stack([x, x], axis=0))\n```\n\nExample:\n```text\ndef body(out_ref):\n  sem_ref = pl.get_global(plgpu.SemaphoreType.REGULAR)\n  block_id = lax.axis_index(\"x\")\n  @pl.when(block_id == 0)\n  def _():\n    pl.semaphore_signal(sem_ref)  # Block 0 signals\n  @pl.when(block_id == 1)\n  def _():\n    pl.semaphore_wait(sem_ref)  # Block 1 waits\n    out_ref[...] = jnp.ones_like(out_ref)\n\nout_shape = jax.ShapeDtypeStruct((128,), jnp.float32)\nplgpu.kernel(body, out_shape=out_shape, grid=(2,), grid_names=(\"x\",))()\n```\n\nExample:\n```text\nref[...] = value\nvalue2 = ref[...]\n```\n\nExample:\n```text\nsmem_ref[...] = value\nplgpu.commit_smem()\nplgpu.copy_smem_to_gmem(smem_ref, ...)\n```\n\nExample:\n```text\nsmem_ref[...] = value\nplgpu.commit_smem()\nplgpu.wgmma(smem_ref, ...)\n```\n\nExample:\n```text\nv = plgpu.load(smem_ref)\nplgpu.commit_smem()\nplgpu.copy_gmem_to_smem(..., smem_ref, ...)\n```\n\nExample:\n```text\nplgpu.barrier_wait(barrier)\n```\n\nExample:\n```text\nplgpu.barrier_arrive(barrier)\n```\n\nExample:\n```text\ntid = jax.lax.axis_index(\"thread\")\nassert queue.shape == (buffering, *item_shape)\nassert produced.shape == consumed.shape == (buffering,)\n\ndef thread0_body(i, _):\n  slot = jax.lax.rem(i, buffering)\n  @pl.when(i >= buffering)\n  def _await_consumed():\n    plgpu.barrier_wait(consumed.at[slot])  # Wait for consumption of the value before overwriting it\n  # Option 1: Compute the next value\n  queue[slot] = produce()\n  plgpu.barrier_arrive(produced.at[slot])  # Signal the value is ready\n  # Option 2: Produce the value through async_copy\n  # plgpu.copy_gmem_to_smem(..., queue.at[slot], barrier=produced.at[slot])\npl.when(tid == 0)(lambda: jax.lax.fori_loop(0, steps, thread0_body, None))\n\ndef thread1_body(i, _):\n  slot = jax.lax.rem(i, buffering)\n  plgpu.barrier_wait(produced.at[slot])  # Wait for the value to be ready\n  consume(queue[slot])  # Load and compute\n  plgpu.barrier_arrive(consumed.at[slot])  # Signal that the value is consumed\npl.when(tid == 1)(lambda: jax.lax.fori_loop(0, steps, thread1_body, None))\n```\n\nExample:\n```text\n@functools.partial(pl.run_scoped, barrier_ref=plgpu.Barrier(orders_tensor_core=True))\ndef barrier_scope(barrier_ref):\n  plgpu.tcgen05_mma(acc_tmem, lhs_ref, rhs_ref, barrier_ref, accumulate=False)\n  plgpu.barrier_wait(barrier_ref)\n  # We can read the result now\n  result = plgpu.async_load_tmem(acc_tmem)\n  ...\n```\n\nExample:\n```text\ndef collective_smem_reuse(x_gmem, x_gmem2, y_gmem, x_smem, local_barrier, cluster_barrier):\n  plgpu.copy_gmem_to_smem(x_gmem, x_smem, local_barrier, collective_axes=\"cluster\")\n  plgpu.barrier_wait(local_barrier)  # x_smem is ready to be used once the local wait completes\n  y_gmem[0] = x_smem[...]\n  plgpu.barrier_arrive(cluster_barrier)\n  plgpu.barrier_wait(cluster_barrier)  # x_smem can only be reused once the cluster barrier completes\n  plgpu.copy_gmem_to_smem(x_gmem2, x_smem, local_barrier, collective_axes=\"cluster\")\n  plgpu.barrier_wait(local_barrier)  # x_smem is ready to be used once the local wait completes\n  y_gmem[1] = x_smem[...]\n```\n\nExample:\n```text\ndef collective_tmem_reuse(acc_tmem, lhs_ref, rhs_ref, mma_barrier, cluster_barrier):\n  leader_block = lax.axis_index(\"cluster\") == 0\n  @pl.when(leader_block)\n  def _do_mma():\n    plgpu.tcgen05_mma(\n        acc_tmem, lhs_ref.at[0], rhs_ref.at[0], mma_barrier,\n        accumulate=False, collective_axis=\"x\",\n    )\n  plgpu.barrier_wait(mma_barrier)\n  do_something(plgpu.async_load_tmem(acc_tmem))\n  plgpu.wait_load_tmem()  # Ensure the load is complete.\n  plgpu.barrier_arrive(cluster_barrier)\n  plgpu.barrier_wait(cluster_barrier)  # acc_tmem can only be reused once the cluster barrier completes\n  @pl.when(leader_block)\n  def _do_mma():\n    plgpu.tcgen05_mma(\n        acc_tmem, lhs_ref.at[1], rhs_ref.at[1], mma_barrier,\n        accumulate=False, collective_axis=\"x\",\n    )\n  ...\n```\n\nExample:\n```text\ndef exchange_shards(x_ref, y_ref, done_sem):\n  other_dev_id = 1 - lax.axis_index(\"x\")  # We assume two devices\n  neighbor_ref = plgpu.remote_ref(y_ref, other_dev_id)\n  neighbor_ref[...] = x_ref[...]  # This will write over NVLINK\n  pl.semaphore_signal(done_sem, device_id=other_dev_id)  # Signal that the write is complete\n  pl.semaphore_wait(done_sem)  # Wait for the other device to write to our memory\n\nmesh = jax.make_mesh((2,), (\"x\",))\ny = jax.jit(\n    jax.shard_map(\n      lambda x: plgpu.kernel(exchange_shards, out_shape=x,\n                             scratch_shapes=[plgpu.SemaphoreType.REGULAR])(x),\n      mesh=mesh, in_specs=P(\"x\"), out_specs=P(\"x\"), check_vma=False,\n    )\n)(x)\n```\n\nExample:\n```text\n@functools.partial(\n    plgpu.kernel,\n    grid=grid,\n    grid_names=grid_names,\n    scratch_shapes=dict(\n        result_ref=plgpu.TryCancelResultRef(),\n        barrier_ref=plgpu.Barrier()\n    )\n)\ndef kernel(result_ref, barrier_ref):\n  plgpu.try_cluster_cancel(result_ref, barrier_ref)\n  # ... do work\n  plgpu.barrier_wait(barrier_ref)\n  grid_idxs, success = plgpu.query_cluster_cancel(result_ref, grid_names)\n```\n\nExample:\n```text\n@plgpu.dynamic_scheduling_loop(\n  grid_names=grid_names,\n  thread_axis=thread_name  # Required if using multiple threads in a kernel.\n)\ndef body(loop_info):\n  grid_indices = loop_info.index\n  # ... do work\n```\n\nExample:\n```text\ndef body(in_gmem_ref, out_gmem_ref, smem_ref, barrier):\n  plgpu.copy_gmem_to_smem(in_gmem_ref, smem_ref, barrier)\n  plgpu.barrier_wait(barrier)\n  ...\n\nplgpu.kernel(\n  body,\n  out_shape=...,\n  scratch_shapes=[plgpu.SMEM(x.shape, x.dtype), plgpu.Barrier()],\n)\n```\n\nExample:\n```text\ndef body(in_gmem_ref, in_gmem_ref2, out_gmem_ref, smem_ref, smem_ref2, barrier):\n  plgpu.copy_gmem_to_smem(in_gmem_ref, smem_ref, barrier)\n  plgpu.copy_gmem_to_smem(in_gmem_ref2, smem_ref2, barrier)\n  plgpu.barrier_wait(barrier)  # Awaits both copies\n  ...\n\nplgpu.kernel(\n  body,\n  out_shape=...,\n  # Barrier is allocated with 2 arrivals.\n  scratch_shapes=[plgpu.SMEM(x.shape, x.dtype), plgpu.Barrier(num_arrivals=2)],\n)\n```\n\nExample:\n```text\ndef body(in_gmem_ref, in_gmem_ref2, out_gmem_ref, smem_ref, smem_ref2, barrier):\n  block_id = lax.axis_index(\"cluster\")\n  # Both blocks in the cluster load the same data into smem_ref, so we can use\n  # a collective copy here.\n  plgpu.copy_gmem_to_smem(in_gmem_ref, smem_ref, barrier, collective_axes=\"cluster\")\n  # Each block in the cluster loads a different slice of in_gmem_ref2, so we\n  # are not allowed to use collective copies.\n  plgpu.copy_gmem_to_smem(in_gmem_ref2.at[block_id], smem_ref2, barrier)\n  plgpu.barrier_wait(barrier)  # Awaits both copies\n  ...\n\nplgpu.kernel(\n  body,\n  out_shape=...,\n  # Barrier is allocated with 2 arrivals.\n  scratch_shapes=[plgpu.SMEM(x.shape, x.dtype), plgpu.Barrier(num_arrivals=2)],\n)\n```\n\nExample:\n```text\nplgpu.copy_gmem_to_smem(\n    x_gmem,  # [K, N]\n    x_smem,  # [K, N]\n    tma_barrier,\n    collective_axes=\"x\",\n    leader_tracked=plgpu.CopyPartition.REPLICATED,\n)\n```\n\nExample:\n```text\ndef copy_out(x_smem, y_smem, x_gmem, y_gmem):\n  plgpu.copy_smem_to_gmem(x_smem, x_gmem)\n  plgpu.copy_smem_to_gmem(y_smem, y_gmem)\n  plgpu.wait_smem_to_gmem(1, wait_read_only=True)\n  # At this point we know that the data of x_smem has been read, but we don't\n  # yet know that x_gmem contains the updated data.\n  plgpu.wait_smem_to_gmem(1)\n  # At this point we know that the x_smem -> x_gmem copy is done, but we know\n  # nothing about the y_smem -> y_gmem copy.\n  plgpu.wait_smem_to_gmem(0)\n  # At this point we know that both copies are complete.\n```\n\nExample:\n```text\ndef copy_out(x_smem, y_smem, x_gmem, y_gmem):\n  plgpu.copy_smem_to_gmem(x_smem, x_gmem, commit_group=False)\n  plgpu.copy_smem_to_gmem(y_smem, y_gmem)  # Implicitly commits both copies\n  plgpu.wait_smem_to_gmem(1)\n  # At this point we only know that no SMEM to GMEM copies other than the two\n  # above are active.\n  plgpu.wait_smem_to_gmem(0)\n  # Only now we know that both copies above have completed.\n```\n\nExample:\n```text\n@functools.partial(\n    self.pallas_call,\n    out_shape=jax.ShapeDtypeStruct(out_shape, dtype),\n    out_specs=plgpu.BlockSpec(memory_space=plgpu.SMEM, transforms=transforms),\n    in_specs=(\n        pl.BlockSpec(memory_space=plgpu.GMEM),\n        pl.BlockSpec(memory_space=plgpu.SMEM),\n    ),\n    scratch_shapes=[plgpu.Barrier()],\n)\ndef kernel(x_ref_gmem, idx_ref, o_ref, barrier_ref):\n  idxs = plgpu.load(idx_ref, layout=plgpu.Layout.TMA_INDICES)\n  plgpu.copy_gmem_to_smem(x_ref_gmem.at[idxs], o_ref, barrier_ref)\n  plgpu.barrier_wait(barrier_ref)\n```\n\nExample:\n```text\ndef exchange_shards(x_ref, y_ref, smem_ref, local_barrier, done_sem):\n  plgpu.copy_gmem_to_smem(x_ref, smem_ref, local_barrier)  # Local copy\n  plgpu.barrier_wait(local_barrier)\n  other_dev_id = 1 - lax.axis_index(\"x\")  # We assume two devices\n  neighbor_ref = plgpu.remote_ref(y_ref, other_dev_id)\n  plgpu.copy_smem_to_gmem(smem_ref, neighbor_ref)\n  plgpu.wait_smem_to_gmem(0)  # Wait for the asynchronous write to complete\n  pl.semaphore_signal(done_sem, device_id=other_dev_id)  # Signal that the write is complete\n  pl.semaphore_wait(done_sem)  # Wait for the other device to write to our memory\n\nmesh = jax.make_mesh((2,), (\"x\",))\ny = jax.jit(\n    jax.shard_map(\n      lambda x: plgpu.kernel(\n          exchange_shards,\n          out_shape=x,\n          scratch_shapes=[x, plgpu.Barrier(), plgpu.Semaphore.REGULAR]\n      )(x),\n      mesh=mesh, in_specs=P(\"x\"), out_specs=P(\"x\"), check_vma=False,\n    )\n)(x)\n```\n\nExample:\n```text\nimport functools\nimport jax\nimport jax.numpy as jnp\nimport torch\n\n@functools.partial(\n    pl.pallas_call, out_shape=jax.ShapeDtypeStruct([128], jnp.int32)\n)\ndef add_kernel(x_ref, y_ref, o_ref):\n  o_ref[...] = x_ref[...] + y_ref[...]\n\nx = torch.arange(128, dtype=torch.int32, device=\"cuda\")\ny = x * x\nout = plgpu.as_torch_kernel(add_kernel)(x, y)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.746Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":44,"totalLines":613,"estimatedTokens":4478}}61{"id":"doc-omnistaging_jax_documentation-230b1336","source":"documentation","title":"Omnistaging — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/4410-omnistaging.html","text":"Example:\n```text\njax.config.disable_omnistaging()\n```\n\nExample:\n```text\n@jit\ndef f(x):\n  input_size = jnp.prod(x.shape)\n  if input_size > 100:\n    ...\n```\n\nExample:\n```text\nimport numpy as np\n\n@jit\ndef f(x):\n  input_size = np.prod(x.shape)\n  if input_size > 100:\n    ...\n```\n\nExample:\n```text\nfrom jax import jit\nimport jax.numpy as jnp\n\n@jit\ndef f(x):\n  y = jnp.add(1, 1)\n  return x * y\n\nf(3)\n```\n\nExample:\n```text\nENTRY jit_f.6 {\n  constant.2 = pred[] constant(false)\n  parameter.1 = s32[] parameter(0)\n  constant.3 = s32[] constant(2)\n  multiply.4 = s32[] multiply(parameter.1, constant.3)\n  ROOT tuple.5 = (s32[]) tuple(multiply.4)\n}\n```\n\nExample:\n```text\nENTRY jit_f.8 {\n  constant.2 = pred[] constant(false)\n  parameter.1 = s32[] parameter(0)\n  constant.3 = s32[] constant(1)\n  constant.4 = s32[] constant(1)\n  add.5 = s32[] add(constant.3, constant.4)\n  multiply.6 = s32[] multiply(parameter.1, add.5)\n  ROOT tuple.7 = (s32[]) tuple(multiply.6)\n}\n```\n\nExample:\n```text\nimport jax.numpy as jnp\nfrom jax import lax\n\n@jit\ndef select_tril(x):\n  mask = jnp.arange(x.shape[0])[:, None] > jnp.arange(x.shape[1])\n  return lax.select(mask, x, jnp.zeros_like(x))  # lax.select is like jnp.where\n\nx = np.arange(12).reshape((3, 4))\nselect_tril(x)\n```\n\nExample:\n```text\nENTRY jit_select_tril.8 {\n  constant.3 = pred[] constant(false)\n  constant.1 = pred[3,4]{1,0} constant({...})\n  parameter.2 = s32[3,4]{1,0} parameter(0)\n  constant.4 = s32[] constant(0)\n  broadcast.5 = s32[3,4]{1,0} broadcast(constant.4), dimensions={}\n  select.6 = s32[3,4]{1,0} select(constant.1, parameter.2, broadcast.5)\n  ROOT tuple.7 = (s32[3,4]{1,0}) tuple(select.6)\n}\n```\n\nExample:\n```text\nENTRY jit_select_tril.16 {\n  constant.4 = pred[] constant(false)\n  iota.1 = s32[3]{0} iota(), iota_dimension=0\n  broadcast.5 = s32[3,1]{1,0} broadcast(iota.1), dimensions={0}\n  reshape.7 = s32[3]{0} reshape(broadcast.5)\n  broadcast.8 = s32[3,4]{1,0} broadcast(reshape.7), dimensions={0}\n  iota.2 = s32[4]{0} iota(), iota_dimension=0\n  broadcast.6 = s32[1,4]{1,0} broadcast(iota.2), dimensions={1}\n  reshape.9 = s32[4]{0} reshape(broadcast.6)\n  broadcast.10 = s32[3,4]{1,0} broadcast(reshape.9), dimensions={1}\n  compare.11 = pred[3,4]{1,0} compare(broadcast.8, broadcast.10), direction=GT\n  parameter.3 = s32[3,4]{1,0} parameter(0)\n  constant.12 = s32[] constant(0)\n  broadcast.13 = s32[3,4]{1,0} broadcast(constant.12), dimensions={}\n  select.14 = s32[3,4]{1,0} select(compare.11, parameter.3, broadcast.13)\n  ROOT tuple.15 = (s32[3,4]{1,0}) tuple(select.14)\n}\n```\n\nExample:\n```text\nfrom jax import jit\nimport jax.numpy as jnp\n\n@jit\ndef ex1(x):\n  size = jnp.prod(jnp.array(x.shape))\n  return x.reshape((size,))\n\nex1(jnp.ones((3, 4)))\n```\n\nExample:\n```text\n[... full traceback ...]\n  File \"/home/mattjj/packages/jax/jax/core.py\", line 862, in raise_concretization_error\n    raise ConcretizationTypeError(msg)\njax.core.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected.\n\nThe error arose in jax.numpy.reshape.\n\nWhile tracing the function ex1 at ex1.py:4, this value became a tracer due to JAX operations on these lines:\n\n  operation c:int32[] = reduce_prod[ axes=(0,) ] b:int32[2]\n    from line ex1.py:6 (ex1)\n\nYou can use transformation parameters such as `static_argnums` for `jit` to avoid tracing particular arguments of transformed functions.\n\nSee https://docs.jax.dev/en/latest/faq.html#abstract-tracer-value-encountered-where-concrete-value-is-expected-error for more information.\n\nEncountered tracer value: Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>\n```\n\nExample:\n```text\nfrom jax import jit\nfrom jax import random\n\nkey = random.PRNGKey(0)\n\ndef init():\n  global key\n  key, subkey = random.split(key)\n  return random.normal(subkey, ())\n\nprint(init())  # -1.2515389\nprint(init())  # -0.58665067\n\ninit = jit(init)\nprint(init())  # 0.48648298\nprint(init())  # 0.48648298  !!\n```\n\nExample:\n```text\nprint(key) # Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace(level=0/1)>\n```\n\nExample:\n```text\nrandom.normal(key, ())\n```\n\nExample:\n```text\n[... full stack trace …]\n  File \"/home/mattjj/packages/jax/jax/interpreters/partial_eval.py\", line 836, in _assert_live\n    raise core.escaped_tracer_error(msg)\njax.core.UnexpectedTracerError: Encountered an unexpected tracer. Perhaps this tracer escaped through global state from a previously traced function.\nThe functions being transformed should not save traced values to global state. Detail: tracer created on line example.py:8 (init).\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.750Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":15,"totalLines":184,"estimatedTokens":1132}}62{"id":"doc-custom_jvp_vjp_rules_for_jax_transformable_funct-93ba5c77","source":"documentation","title":"Custom JVP/VJP rules for JAX-transformable functions — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/2026-custom-derivatives.html","text":"Example:\n```text\n# old custom_transforms api to be replaced\n@jax.custom_transforms\ndef f(x):\n  return 2. * x\n\n# f_vjp :: a -> (b, CT b --o CT a)\ndef f_vjp(x):\n  return f(x), lambda g: 3. * x  # 3 instead of 2\n\njax.defvjp_all(f, f_vjp)\n\ngrad(f)(1.)  # 3.\nvmap(grad(f))(np.ones(4))  # [3., 3., 3., 3.]\ngrad(lambda x: vmap(f)(x).sum())(np.ones(4))  # [2., 2., 2., 2.]\n```\n\nExample:\n```text\n{ lambda  ; ; a.\n  let b = f_primitive a\n  in [b] }\n```\n\nExample:\n```text\n{ lambda  ; ; a.\n  let b = mul 2. a\n  in [b] }\n```\n\nExample:\n```text\nvmap(f)(xs) == np.stack([f(x) for x in xs])\n```\n\nExample:\n```text\njvp(vmap(f))(xs) == jvp(lambda xs: np.stack([f(x) for x in xs]))\n```\n\nExample:\n```text\n# old custom_transforms api to be replaced\n@jax.custom_transforms\ndef f(x):\n  if x > 0:\n    return x\n  else:\n    return 0.\n\ndef f_vjp(x):\n  return ...\n\njax.defvjp_all(f, f_vjp)\n\ngrad(f)(1.)  # Error!\n```\n\nExample:\n```text\nvmap(call(f)) == call(vmap(f))\n```\n\nExample:\n```text\nvmap(custom_jvp_call(f, f_jvp)) == custom_jvp_call(vmap(f), vmap(f_jvp))\n```\n\nExample:\n```text\njvp(call(f)) == call(jvp(f))\n\njvp(custom_jvp_call(f, f_jvp)) == f_jvp\n```\n\nExample:\n```text\neval(call(f)) == eval(f)\njit(call(f)) == hlo_call(jit(f))\n\neval(custom_jvp_call(f, f_jvp)) == eval(f)\njit(custom_jvp_call(f, f_jvp)) == hlo_call(jit(f))\n```\n\nExample:\n```text\n# f :: a -> b\n@jax.custom_jvp\ndef f(x):\n  return np.sin(x)\n\n# f_jvp :: (a, T a) -> (b, T b)\ndef f_jvp(primals, tangents):\n  x, = primals\n  t, = tangents\n  return f(x), np.cos(x) * t\n\nf.defjvp(f_jvp)\n```\n\nExample:\n```text\n# f :: a -> b\n@jax.custom_vjp\ndef f(x):\n  return np.sin(x)\n\n# f_fwd :: a -> (b, c)\ndef f_fwd(x):\n  return f(x), np.cos(x)\n\n# f_bwd :: (c, CT b) -> CT a\ndef f_bwd(cos_x, g):\n  return (cos_x * g,)\n\nf.defvjp(f_fwd, f_bwd)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.751Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":12,"totalLines":121,"estimatedTokens":445}}63{"id":"doc-jax_prng_design_jax_documentation-b117d8d1","source":"documentation","title":"JAX PRNG Design — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/263-prng.html","text":"Example:\n```text\ndef foo(): return bar() + baz()\ndef bar(): return rand(RNG, (3, 4))\ndef baz(): return rand(RNG, (3, 4))\ndef main():\n  global RNG\n  RNG = RandomState(0)\n  return foo()\n```\n\nExample:\n```text\nIn [1]: rng = np.random.RandomState(0)\n\nIn [2]: rng.randn(2)\nOut[2]: array([1.76405235, 0.40015721])\n\nIn [3]: rng = np.random.RandomState(0)\n\nIn [4]: np.stack([rng.randn() for _ in range(2)])\nOut[4]: array([1.76405235, 0.40015721])\n```\n\nExample:\n```text\ndef foo(rng_1):\n   y, rng_2 = baz(rng_1)\n   z, rng_3 = bar(rng_2)\n   return y + z, rng_3\n\ndef bar(x, rng):\n  val, new_rng = rand(rng, (3, 4))\n  return val, new_rng\n\ndef baz(x, rng):\n  val, new_rng = rand(rng, (3, 4))\n  return val, new_rng\n\ndef main():\n  foo(RandomState(0))\n```\n\nExample:\n```text\ndef foo(rng_1):\n   rng_2, rng_3 = split(rng_1, 2)\n   return bar(rng_2) + baz(rng_3)\n\ndef bar(x, rng):\n  return rand(rng, (3, 4))\n\ndef baz(x, rng):\n  return rand(rng, (3, 4))\n\ndef main():\n  foo(RandomState(0))\n```\n\nExample:\n```text\ntype Sample = Int256\ntype Key = Sample  -- important identification for splitting\ntype Count = Int32\n\nhash :: Key -> Count -> Int256  -- output type equal to Key and Sample\n\nsplit :: Key -> (Key, Key)\nsplit key = (hash key 0, hash key 1)\n\ndraw_samples :: Key -> Int -> [Sample]\ndraw_samples key n = map (hash key) [1..n]\n```\n\nExample:\n```text\nrng = lax.rng.new_rng()\nfor i in xrange(num_steps):\n  rng, rng_input = lax.rng.split(rng)\n  params = compiled_update(rng_input, params, next(batches))\n```\n\nExample:\n```text\ndef Dropout(rate, mode='train'):\n  def init_fun(input_shape):\n    return input_shape, ()\n  def apply_fun(rng, params, inputs):\n    if mode == 'train':\n      keep = lax.random.bernoulli(rng, rate, inputs.shape)\n      return np.where(keep, inputs / rate, 0)\n    else:\n      return inputs\n  return init_fun, apply_fun\n```\n\nExample:\n```text\ndef serial(*layers):\n  init_funs, apply_funs = zip(*layers)\n  def init_fun(input_shape):\n    ...\n  def apply_fun(rng, params, inputs):\n    rngs = split(rng, len(layers))\n    for rng, param, apply_fun in zip(rngs, params, apply_funs):\n      inputs = apply_fun(rng, param, inputs)\n    return inputs\n  return init_fun, apply_fun\n\ndef parallel(*layers):\n  init_funs, apply_funs = zip(*layers)\n  def init_fun(input_shape):\n    ...\n  def apply_fun(rng, params, inputs):\n    rngs = split(rng, len(layers))\n    return [f(r, p, x) for f, r, p, x in zip(apply_funs, rngs, params, inputs)]\n  return init_fun, apply_fun\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.752Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":8,"totalLines":120,"estimatedTokens":617}}64{"id":"doc-autodidax2_part_1_jax_from_scratch_again_jax_doc-5c2ded2d","source":"documentation","title":"Autodidax2, part 1: JAX from scratch, again — JAX documentation","url":"https://docs.jax.dev/en/latest/autodidax2_part1.html","text":"Example:\n```text\ndef foo(x):\n  return mul(x, add(x, 3.0))\n```\n\nExample:\n```text\nfrom enum import Enum, auto\nfrom contextlib import contextmanager\nfrom typing import Any\n\n# The full (closed) set of primitive operations\nclass Op(Enum):\n  add = auto()  # addition on floats\n  mul = auto()  # multiplication on floats\n\n# Interpreters have rules for handling each primitive operation.\nclass Interpreter:\n  def interpret_op(self, op: Op, args: tuple[Any, ...]):\n    assert False, \"subclass should implement this\"\n\n# Our first interpreter is the \"evaluating interpreter\" which performs ordinary\n# concrete evaluation.\nclass EvalInterpreter:\n  def interpret_op(self, op, args):\n    assert all(isinstance(arg, float) for arg in args)\n    match op:\n      case Op.add:\n        x, y = args\n        return x + y\n      case Op.mul:\n        x, y = args\n        return x * y\n      case _:\n        raise ValueError(f\"Unrecognized primitive op: {op}\")\n\n# The current interpreter is initially the evaluating interpreter.\ncurrent_interpreter = EvalInterpreter()\n\n# A context manager for temporarily changing the current interpreter\n@contextmanager\ndef set_interpreter(new_interpreter):\n  global current_interpreter\n  prev_interpreter = current_interpreter\n  try:\n    current_interpreter = new_interpreter\n    yield\n  finally:\n    current_interpreter = prev_interpreter\n\n# The user-facing functions `mul` and `add` dispatch to the current interpreter.\ndef add(x, y): return current_interpreter.interpret_op(Op.add, (x, y))\ndef mul(x, y): return current_interpreter.interpret_op(Op.mul, (x, y))\n```\n\nExample:\n```text\nprint(foo(2.0))\n```\n\nExample:\n```text\n10.0\n```\n\nExample:\n```text\nprint((foo(2.00001) - foo(2.0)) / 0.00001)\n```\n\nExample:\n```text\n7.000009999913458\n```\n\nExample:\n```text\nprint(foo(2.00001))\n```\n\nExample:\n```text\n10.0000700001\n```\n\nExample:\n```text\nx + y = (xp + xt * eps) + (yp + yt * eps)\n       =   (xp + yp)             # primal component\n         + (xt + yt) * eps       # tangent component\n```\n\nExample:\n```text\nx * y = (xp + xt * eps) * (yp + yt * eps)\n       =    (xp * yp)                        # primal component\n          + (xp * yt + xt * yp) * eps        # tangent component\n          + (xt * yt)           * eps * eps  # quadratic component, vanishes in the eps->0 limit\n```\n\nExample:\n```text\nfrom dataclasses import dataclass\n\n# A primal-tangent pair is conventionally called a \"dual number\"\n@dataclass\nclass DualNumber:\n  primal  : float\n  tangent : float\n\ndef add_dual(x : DualNumber, y: DualNumber) -> DualNumber:\n  return DualNumber(x.primal + y.primal, x.tangent + y.tangent)\n\ndef mul_dual(x : DualNumber, y: DualNumber) -> DualNumber:\n  return DualNumber(x.primal * y.primal, x.primal * y.tangent + x.tangent * y.primal)\n\ndef foo_dual(x : DualNumber) -> DualNumber:\n  return mul_dual(x, add_dual(x, DualNumber(3.0, 0.0)))\n\nprint (foo_dual(DualNumber(2.0, 1.0)))\n```\n\nExample:\n```text\nDualNumber(primal=10.0, tangent=7.0)\n```\n\nExample:\n```text\n# This is like DualNumber above except that is also has a pointer to the\n# interpreter it belongs to, which is needed to avoid \"perturbation confusion\"\n# in higher order differentiation.\n@dataclass\nclass TaggedDualNumber:\n  interpreter : Interpreter\n  primal  : float\n  tangent : float\n\nclass JVPInterpreter(Interpreter):\n  def __init__(self, prev_interpreter: Interpreter):\n    # We keep a pointer to the interpreter that was current when this\n    # interpreter was first invoked. That's the context in which our\n    # rules should run.\n    self.prev_interpreter = prev_interpreter\n\n  def interpret_op(self, op, args):\n    args = tuple(self.lift(arg) for arg in args)\n    with set_interpreter(self.prev_interpreter):\n      match op:\n        case Op.add:\n          # Notice that we use `add` and `mul` here, which are the\n          # interpreter-dispatching functions defined earlier.\n          x, y = args\n          return self.dual_number(\n              add(x.primal, y.primal),\n              add(x.tangent, y.tangent))\n\n        case Op.mul:\n          x, y = args\n          x = self.lift(x)\n          y = self.lift(y)\n          return self.dual_number(\n              mul(x.primal, y.primal),\n              add(mul(x.primal, y.tangent), mul(x.tangent, y.primal)))\n\n  def dual_number(self, primal, tangent):\n    return TaggedDualNumber(self, primal, tangent)\n\n  # Lift a constant value (constant with respect to this interpreter) to\n  # a TaggedDualNumber.\n  def lift(self, x):\n    if isinstance(x, TaggedDualNumber) and x.interpreter is self:\n      return x\n    else:\n      return self.dual_number(x, 0.0)\n\ndef jvp(f, primal, tangent):\n  jvp_interpreter = JVPInterpreter(current_interpreter)\n  dual_number_in = jvp_interpreter.dual_number(primal, tangent)\n  with set_interpreter(jvp_interpreter):\n    result = f(dual_number_in)\n  dual_number_out = jvp_interpreter.lift(result)\n  return dual_number_out.primal, dual_number_out.tangent\n\n# Let's try it out:\nprint(jvp(foo, 2.0, 1.0))\n\n# Because we were careful to consider nesting interpreters, higher-order AD\n# works out of the box:\n\ndef derivative(f, x):\n  _, tangent = jvp(f, x, 1.0)\n  return tangent\n\ndef nth_order_derivative(n, f, x):\n  if n == 0:\n    return f(x)\n  else:\n    return derivative(lambda x: nth_order_derivative(n-1, f, x), x)\n```\n\nExample:\n```text\n(10.0, 7.0)\n```\n\nExample:\n```text\nprint(nth_order_derivative(0, foo, 2.0))\n```\n\nExample:\n```text\nprint(nth_order_derivative(1, foo, 2.0))\n```\n\nExample:\n```text\n7.0\n```\n\nExample:\n```text\nprint(nth_order_derivative(2, foo, 2.0))\n```\n\nExample:\n```text\n2.0\n```\n\nExample:\n```text\n# The rest are zero because `foo` is only a second-order polymonial\nprint(nth_order_derivative(3, foo, 2.0))\n```\n\nExample:\n```text\n0.0\n```\n\nExample:\n```text\nprint(nth_order_derivative(4, foo, 2.0))\n```\n\nExample:\n```text\ndef f(x):\n  # g is constant in its (ignored) argument `y`. Its derivative should be zero\n  # but our AD will mess it up if we don't distinguish perturbations from\n  # different interpreters.\n  def g(y):\n    return x\n  should_be_zero = derivative(g, 0.0)\n  return mul(x, should_be_zero)\n\nprint(derivative(f, 0.0))\n```\n\nExample:\n```text\nVar = str           # Variables are just strings in this untyped IR\nAtom = Var | float  # Atoms (arguments to operations) can be variables or (float) literals\n\n# Equation - a single line in our IR like `z = mul(x, y)`\n@dataclass\nclass Equation:\n  var  : Var         # The variable name of the result\n  op   : Op          # The primitive operation we're applying\n  args : tuple[Atom] # The arguments we're applying the primitive operation to\n\n# We call an IR function a \"Jaxpr\", for \"JAX expression\"\n@dataclass\nclass Jaxpr:\n  parameters : list[Var]      # The function's formal parameters (arguments)\n  equations  : list[Equation] # The body of the function, a list of instructions/equations\n  return_val : Atom           # The function's return value\n\n  def __str__(self):\n    lines = []\n    lines.append(', '.join(b for b in self.parameters) + ' ->')\n    for eqn in self.equations:\n      args_str = ', '.join(str(arg) for arg in eqn.args)\n      lines.append(f'  {eqn.var} = {eqn.op}({args_str})')\n    lines.append(self.return_val)\n    return '\\n'.join(lines)\n```\n\nExample:\n```text\nclass StagingInterpreter(Interpreter):\n  def __init__(self):\n    self.equations = []         # A mutable list of all the ops we've seen so far\n    self.name_counter = 0  # Counter for generating unique names\n\n  def fresh_var(self):\n    self.name_counter += 1\n    return \"v_\" + str(self.name_counter)\n\n  def interpret_op(self, op, args):\n    binder = self.fresh_var()\n    self.equations.append(Equation(binder, op, args))\n    return binder\n\ndef build_jaxpr(f, num_args):\n  interpreter = StagingInterpreter()\n  parameters = tuple(interpreter.fresh_var() for _ in range(num_args))\n  with set_interpreter(interpreter):\n    result = f(*parameters)\n  return Jaxpr(parameters, interpreter.equations, result)\n```\n\nExample:\n```text\nprint(build_jaxpr(foo, 1))\n```\n\nExample:\n```text\nv_1 ->\n  v_2 = Op.add(v_1, 3.0)\n  v_3 = Op.mul(v_1, v_2)\nv_3\n```\n\nExample:\n```text\ndef eval_jaxpr(jaxpr, args):\n  # An environment mapping variables to values\n  env = dict(zip(jaxpr.parameters, args))\n  def eval_atom(x): return env[x] if isinstance(x, Var) else x\n  for eqn in jaxpr.equations:\n    args = tuple(eval_atom(x) for x in eqn.args)\n    env[eqn.var] = current_interpreter.interpret_op(eqn.op, args)\n  return eval_atom(jaxpr.return_val)\n\nprint(eval_jaxpr(build_jaxpr(foo, 1), (2.0,)))\n```\n\nExample:\n```text\nprint(jvp(lambda x: eval_jaxpr(build_jaxpr(foo, 1), (x,)), 2.0, 1.0))\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.753Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":29,"totalLines":348,"estimatedTokens":2141}}65{"id":"doc-software_pipelining_jax_documentation-76c2dad0","source":"documentation","title":"Software Pipelining — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/pipelining.html","text":"Example:\n```text\nimport jax\nfrom jax import numpy as jnp\nfrom jax.experimental import pallas as pl\nimport numpy as np\n```\n\nExample:\n```text\n# Note: This is a TPU example.\n\ndef add_matrices_kernel(x_sram_ref, y_sram_ref, z_sram_ref):\n  # Load x and y from SRAM into registers\n  x_regs = x_sram_ref[:, :]\n  y_regs = y_sram_ref[:, :]\n  # Execute a vectorized add\n  z_regs = x_regs + y_regs\n  # Store the output values in registers back into SRAM\n  z_sram_ref[:, :] = z_regs\n\n\ndef add_matrices(x: jax.Array, y: jax.Array) -> jax.Array:\n  # pallas_call will first allocate scratch buffers for `x` and `y` in SRAM.\n  # It will then copy `x` and `y` from HBM into SRAM.\n  z = pl.pallas_call(\n      add_matrices_kernel, out_shape=jax.ShapeDtypeStruct.like(x)\n  )(x, y)\n  # pallas_call will also copy the output from SRAM back into HBM.\n  return z\n\n\nx, y = jnp.ones((512, 512)), jnp.ones((512, 512))\nadd_matrices(x, y)\n```\n\nExample:\n```text\nArray([[2., 2., 2., ..., 2., 2., 2.],\n       [2., 2., 2., ..., 2., 2., 2.],\n       [2., 2., 2., ..., 2., 2., 2.],\n       ...,\n       [2., 2., 2., ..., 2., 2., 2.],\n       [2., 2., 2., ..., 2., 2., 2.],\n       [2., 2., 2., ..., 2., 2., 2.]], dtype=float32)\n```\n\nExample:\n```text\nfor i in range(N):\n  copy_in(A[i], X)\n  Y = X + 1\n  copy_out(Y, A[i])\n```\n\nExample:\n```text\n# Itr 1\n  copy_in_start(A[0], X)\n  copy_in_wait(X)\n  Y = X + 1\n  copy_out_start(Y, A[0])\n  copy_out_wait(Y)\n\n  # Itr 2\n  copy_in_start(A[1], X)\n  copy_in_wait(X)\n  Y = X + 1\n  copy_out_start(Y, A[1])\n  copy_out_wait(Y)\n\n  # Itr 3\n  copy_in_start(A[2], X)\n  copy_in_wait(X)\n  Y = X + 1\n  copy_out_start(Y, A[2])\n  copy_out_wait(Y)\n\n  # Itr 4\n  copy_in_start(A[3], X)\n  copy_in_wait(X)\n  Y = X + 1\n  copy_out_start(Y, A[3])\n  copy_out_wait(Y)\n```\n\nExample:\n```text\n# Prologue\n  copy_in_start(A[0], X[0])\n  \n  # Itr 1\n  copy_in_start(A[1], X[1])\n  copy_in_wait(X[0])\n  Y[0] = X[0] + 1\n  copy_out_start(Y[0], A[0])\n  copy_out_wait(Y[0])\n\n  # Itr 2 - Steady state\n  copy_in_start(A[2], X[0])\n  copy_in_wait(X[1])\n  Y[1] = X[1] + 1\n  copy_out_start(Y[1], A[1])\n  copy_out_wait(Y[1])\n\n  # Itr 3 - Steady state\n  copy_in_start(A[3], X[1])\n  copy_in_wait(X[0])\n  Y[0] = X[0] + 1\n  copy_out_start(Y[0], A[2])\n  copy_out_wait(Y[0])\n\n  # Itr 4 - No copy-in\n  copy_in_wait(X[1])\n  Y[1] = X[1] + 1\n  copy_out_start(Y[1], A[3])\n  copy_out_wait(Y[1])\n```\n\nExample:\n```text\n# Prologue\n  copy_in_start(A[0], X[0])\n  \n  # Itr 1\n  copy_in_start(A[1], X[1])\n  copy_in_wait(X[0])\n  Y[0] = X[0] + 1\n  copy_out_start(Y[0], A[0])\n\n  # Itr 2 - Steady state\n  copy_in_start(A[2], X[0])\n  copy_in_wait(X[1])\n  Y[1] = X[1] + 1\n  copy_out_start(Y[1], A[1])\n  copy_out_wait(Y[0])\n\n  # Itr 3 - Steady state\n  copy_in_start(A[3], X[1])\n  copy_in_wait(X[0])\n  Y[0] = X[0] + 1\n  copy_out_start(Y[0], A[2])\n  copy_out_wait(Y[1])\n\n  # Itr 4 - No copy-in\n  copy_in_wait(X[1])\n  Y[1] = X[1] + 1\n  copy_out_start(Y[1], A[3])\n  copy_out_wait(Y[0])\n\n  # Epilogue\n  copy_out_wait(Y[1])\n```\n\nExample:\n```text\n# Prologue\ncopy_in_start(A[0], X[0])\n\n# Main loop\nfor i in range(N):\n  cur_slot = i % 2\n  next_slot = (i + 1) % 2\n\n  if i+1 < N:\n    copy_in_start(A[i+1], X[next_slot])\n  \n  copy_in_wait(X[cur_slot])\n  Y[cur_slot] = X[cur_slot] + 1\n  copy_out_start(Y[cur_slot], A[i])\n\n  if i > 0:\n    copy_out_wait(Y[next_slot])\n\n# Epilogue\ncopy_out_wait(Y[1])\n```\n\nExample:\n```text\ndef double_buffered_pipeline(\n    grid: tuple[int, ...],\n    kernel: Callable,\n    in_slices: Callable,\n    out_slices: Callable):\n  # Prologue\n  copy_in_start(in_hbm[in_slices(0)], in_sram[0])\n\n  # Main loop\n  grid_size = prod(grid)\n  for i in range(grid_size):\n    cur_slot = i % 2\n    next_slot = (i + 1) % 2\n    if (i + 1) < grid_size:\n      copy_in_start(in_hbm[in_slices(i+1)], in_sram[next_slot])\n    copy_in_wait(in_sram[cur_slot])\n\n    kernel(in_sram[cur_slot], out_sram[cur_slot])\n\n    copy_out_start(out_sram[cur_slot], out_hbm[out_slices(i)])\n    if i > 0:\n      copy_out_wait(out_sram[next_slot])\n\n  # Epilogue\n  last_slot = (grid_size - 1) % 2\n  copy_out_wait(out_sram[last_slot])\n```\n\nExample:\n```text\n# For grid (N, M, K)\nfor n in range (N):\n  for m in range(M):\n    for k in range(K):\n      kernel()\n```\n\nExample:\n```text\npl.BlockSpec(\n  block_shape: tuple[int, ...],\n  index_map: Callable,\n  memory_space: pl.MemorySpace\n)\n```\n\nExample:\n```text\ndef kernel(*input_buffers, *output_buffers):\n  # ... perform compute\n  # ... store result into output buffers\n```\n\nExample:\n```text\ndef pallas_call(\n  kernel,\n  grid: tuple[int, ...],\n  in_specs: Sequence[PyTree[BlockSpec]],\n  out_specs: PyTree[BlockSpec],\n  out_shape: PyTree[jax.ShapeDtypeStruct],\n) -> Callable:\n```\n\nExample:\n```text\n# Note: This is a TPU example.\n\ntotal_shape = (4096, 4096)\nblock_shape = (512, 512)\n\ndef add_matrices_pipelined_kernel(x_ref, y_ref, o_ref):\n  o_ref[...] = x_ref[...] + y_ref[...]\n\ndef add_matrices_pipelined(x: jax.Array, y: jax.Array):\n  return pl.pallas_call(\n    add_matrices_pipelined_kernel,\n    grid=tuple(total // block for (total, block) in zip(total_shape, block_shape)),\n    in_specs=[\n      pl.BlockSpec(block_shape, index_map=lambda i, j: (i, j)),\n      pl.BlockSpec(block_shape, index_map=lambda i, j: (i, j))\n    ],\n    out_specs=pl.BlockSpec(block_shape, index_map=lambda i, j: (i, j)),\n    out_shape=jax.ShapeDtypeStruct(total_shape, dtype=jnp.float32),\n  )(x, y)\n\nx = jax.random.uniform(jax.random.key(0), total_shape, dtype=jnp.float32)\ny = jax.random.uniform(jax.random.key(1), total_shape, dtype=jnp.float32)\nresult = add_matrices_pipelined(x, y)\nnp.testing.assert_array_equal(\n    result, x + y\n)\n```\n\nExample:\n```text\ndef add_matrices_pipelined_param(\n    x: jax.Array, y: jax.Array, *, bm: int = 256, bn: int = 256\n) -> jax.Array:\n  m, n = x.shape\n  block_spec = pl.BlockSpec((bm, bn), lambda i, j: (i, j))\n  return pl.pallas_call(\n      add_matrices_kernel,\n      out_shape=x,\n      in_specs=[block_spec, block_spec],\n      out_specs=block_spec,\n      grid=(m // bm, n // bn),\n  )(x, y)\n\nnp.testing.assert_array_equal(\n    add_matrices_pipelined_param(x, y, bm=256, bn=256), x + y\n)\nnp.testing.assert_array_equal(\n    add_matrices_pipelined_param(x, y, bm=128, bn=128), x + y\n)\nnp.testing.assert_array_equal(\n    add_matrices_pipelined_param(x, y, bm=512, bn=512), x + y\n)\n```\n\nExample:\n```text\nx = jnp.ones((8, 1024, 1024))\njnp.sum(x, axis=0)\n```\n\nExample:\n```text\nArray([[8., 8., 8., ..., 8., 8., 8.],\n       [8., 8., 8., ..., 8., 8., 8.],\n       [8., 8., 8., ..., 8., 8., 8.],\n       ...,\n       [8., 8., 8., ..., 8., 8., 8.],\n       [8., 8., 8., ..., 8., 8., 8.],\n       [8., 8., 8., ..., 8., 8., 8.]], dtype=float32)\n```\n\nExample:\n```text\n# Note: This is a TPU example.\n\n# Warning: this implementation is incorrect!\ndef incorrect_sum_kernel(x_ref, o_ref):\n  o_ref[...] += x_ref[...]\n\ndef incorrect_sum(x: jax.Array,\n              block_size: tuple[int, ...] = (256, 256)) -> jax.Array:\n  reduction_size, *out_shape = x.shape\n  grid = (reduction_size, *(out // blk for out, blk in zip(out_shape, block_size)))\n  return pl.pallas_call(\n      incorrect_sum_kernel,\n      grid=grid,\n      # None in `block_shape` means we pick a size of 1 and squeeze it away\n      in_specs=[pl.BlockSpec((None, *block_size), lambda i, j, k: (i, j, k))],\n      out_specs=pl.BlockSpec(block_size, lambda i, j, k: (j, k)),\n      out_shape=jax.ShapeDtypeStruct(out_shape, x.dtype),\n  )(x)\n\nresult = incorrect_sum(x)\nprint(result)\n```\n\nExample:\n```text\n[[65. 65. 65. ... 66. 66. 66.]\n [65. 65. 65. ... 66. 66. 66.]\n [65. 65. 65. ... 66. 66. 66.]\n ...\n [71. 71. 71. ... 72. 72. 72.]\n [71. 71. 71. ... 72. 72. 72.]\n [71. 71. 71. ... 72. 72. 72.]]\n```\n\nExample:\n```text\n# Note: This is a TPU example.\n\ndef correct_sum_kernel(x_ref, o_ref):\n  @pl.when(pl.program_id(2) == 0)\n  def _():\n    o_ref[...] = jnp.zeros_like(o_ref)\n  o_ref[...] += x_ref[...]\n\ndef correct_sum(x: jax.Array,\n              block_size: tuple[int, ...] = (256, 256)) -> jax.Array:\n  reduction_size, *out_shape = x.shape\n  # We moved the reduction to the last axis of the grid.\n  grid = (*(out // blk for out, blk in zip(out_shape, block_size)), reduction_size)\n  return pl.pallas_call(\n      correct_sum_kernel,\n      grid=grid,\n      # None in `block_shape` means we pick a size of 1 and squeeze it away\n      in_specs=[pl.BlockSpec((None, *block_size), lambda i, j, k: (k, i, j))],\n      out_specs=pl.BlockSpec(block_size, lambda i, j, k: (i, j)),\n      out_shape=jax.ShapeDtypeStruct(out_shape, x.dtype),\n  )(x)\n\nresult = correct_sum(x)\nprint(result)\n```\n\nExample:\n```text\n[[8. 8. 8. ... 8. 8. 8.]\n [8. 8. 8. ... 8. 8. 8.]\n [8. 8. 8. ... 8. 8. 8.]\n ...\n [8. 8. 8. ... 8. 8. 8.]\n [8. 8. 8. ... 8. 8. 8.]\n [8. 8. 8. ... 8. 8. 8.]]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.756Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":21,"totalLines":393,"estimatedTokens":2164}}66{"id":"doc-pallas_design_jax_documentation-9e72f7af","source":"documentation","title":"Pallas Design — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/design/design.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.experimental import pallas as pl\n\ndef add_kernel(x_ref, y_ref, o_ref):\n  # In this code, `x_ref`, `y_ref` and `o_ref` are (8,)-shaped `Ref`s\n  x = x_ref[:]\n  y = y_ref[:]\n  o_ref[:] = x + y\nx, y = jnp.arange(8), jnp.arange(8, 16)\nadd = pl.pallas_call(add_kernel, out_shape=jax.ShapeDtypeStruct((8,), jnp.int32))\nadd(x, y)\n```\n\nExample:\n```text\ndef f(x_ref, o_ref):\n  # Using vanilla Python indexing\n  x = x_ref[0, 2:5, :]\n  # Or via Numpy advanced int indexing\n  o_ref[jnp.arange(3), :] = x\n\n# Note that in order to use NumPy advanced int indexing, you need to broadcast the indices against each other into the desired multidimensional shape:\ndef f(x_ref):\n  # Assume x_ref is (8, 4) and we want to read out a (2, 3) slice\n  x = x_ref[jnp.arange(2)[..., None], jnp.arange(3)[None, ...]]\n```\n\nExample:\n```text\ndef f(x_ref, o_ref):\n  # Reading from memory via pallas.load\n  x = pl.load(x_ref, (0, slice(2, 5), slice(None)))\n  # Using integer indexing automatically broadcasts\n  x = pl.load(x_ref, (0, 2 + jnp.arange(3), slice(None)))\n  # You can also use `pl.dynamic_slice` (`pl.ds` for short) objects as well\n  pl.store(o_ref, (0, pl.ds(start=2, size=3), slice(None)), x)\n```\n\nExample:\n```text\ndef f(x_ref, o_ref):\n  # Reading from memory via pallas.load\n  idx = jnp.arange(8)\n  mask = idx < 5\n  x = pl.load(x_ref, (idx,), mask=mask, other=float('-inf'))\n```\n\nExample:\n```text\ndef f(x_ref, o_ref):\n  i = pl.program_id(axis=0)  # execution index in the first axis of the grid\n  o_ref[i] = jnp.exp(x_ref[i])\n```\n\nExample:\n```text\ndef pallas_call(\n    kernel: Callable,\n    out_shape: Sequence[jax.ShapeDtypeStruct],\n    *,\n    in_specs: Sequence[Spec],\n    out_specs: Sequence[Spec],\n    grid: Optional[Tuple[int, ...]] = None) -> Callable:\n  ...\n```\n\nExample:\n```text\ndef pallas_call(kernel, out_shape, *, in_specs, out_specs, grid):\n  def execute(*args):\n    outputs = map(empty_ref, out_shape)\n    grid_indices = map(range, grid)\n    for indices in itertools.product(*grid_indices): # Could run in parallel!\n      local_inputs = [in_spec.transform(arg, indices) for arg, in_spec in\n                      zip(args, in_specs)]\n      local_outputs = [out_spec.transform(arg, indices) for arg, out_spec  in\n                       zip(outputs, out_specs)]\n      kernel(*local_inputs, *local_outputs) # writes to outputs\n  return execute\n```\n\nExample:\n```text\nclass BlockSpec:\n  index_map: Callable[[Tuple[Int, ...]], Tuple[Int, ...]]\n  block_shape: Tuple[Optional[int], ...]\n\n  def transform(self, ref, *loop_indices):\n    block_indices = self.transform_function(loop_indices)\n    # Returns a view of `ref` starting at `block_indices` of shape self.block_shape\n    ...\n```\n\nExample:\n```text\ndef make_kernel(eltwise_kernel):\n  def add(x_ref, y_ref, o_ref):\n    x = pl.load(x_ref, ())\n    y = pl.load(y_ref, ())\n    pl.store(o_ref, (), eltwise_kernel(x + y))\n  return add\n\nkernel1 = make_kernel(lambda x: x * 2)\nkernel2 = make_kernel(jnp.exp)\n\npl.pallas_call(kernel1, out_shape=x, grid=1)(1., 1.)\npl.pallas_call(kernel2, out_shape=x, grid=1)(1., 1.)\n```\n\nExample:\n```text\ndef add_kernel(x_ref, y_ref, o_ref):\n  # In this code, `x_ref`, `y_ref` and `o_ref` are (2,)-shaped `Ref`s\n  x = x_ref[:]\n  y = y_ref[:]\n  o_ref[:] = x + y\nx, y = jnp.arange(8), jnp.arange(8, 16)\nadd = pl.pallas_call(\n    add_kernel,\n    out_shape=jax.ShapeDtypeStruct((8,), jnp.int32),\n    in_specs=[\n        pl.BlockSpec((2,), lambda i: i),\n        pl.BlockSpec((2,), lambda i: i)\n    ],\n    out_specs=pl.BlockSpec((2,), lambda i: i),\n    grid=(4,))\nadd(x, y)\n```\n\nExample:\n```text\ndef matmul_kernel(x_ref, y_ref, o_ref, *, activation, block_k):\n  acc = jnp.zeros((x_ref.shape[0], y_ref.shape[1]), jnp.float32)\n  for k in range(x_ref.shape[1] // block_k):\n    x = x_ref[:, k*block_k:(k+1)*block_k]\n    y = y_ref[k*block_k:(k+1)*block_k, :]\n    acc += x @ y\n  o_ref[:, :] = activation(acc).astype(o_ref.dtype)\n\nx, y = jnp.ones((512, 256)), jnp.ones((256, 1024))\nblock_shape = 128, 256, 128\n\n@jax.jit(static_argnames=[\"block_shape\", \"activation\"])\ndef matmul(x, y, *, block_shape, activation):\n  block_m, block_n, block_k = block_shape\n  fused_matmul = pl.pallas_call(\n      partial(matmul_kernel, block_k=block_k, activation=activation),\n      out_shape=jax.ShapeDtypeStruct((x.shape[0], y.shape[1],), jnp.float32),\n      in_specs=[\n          pl.BlockSpec((block_m, x.shape[1]), lambda i, j: (i, 0)),\n          pl.BlockSpec((y.shape[0], block_n), lambda i, j: (0, j))\n      ],\n      out_specs=pl.BlockSpec((block_m, block_n), lambda i, j: (i, j)),\n      grid=(4, 4),\n  )\n  return fused_matmul(x, y)\n\nz = matmul(x, y, block_shape=block_shape, activation=jax.nn.gelu)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.757Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":11,"totalLines":164,"estimatedTokens":1178}}67{"id":"doc-building_from_source_jax_documentation-34fbc7b0","source":"documentation","title":"Building from source — JAX documentation","url":"https://docs.jax.dev/en/latest/developer.html","text":"Example:\n```text\ngit clone https://github.com/jax-ml/jax\ncd jax\n```\n\nExample:\n```text\npip install jaxlib\n```\n\nExample:\n```text\npython build/build.py build --wheels=jaxlib --verbose\npip install dist/*.whl  # installs jaxlib (includes XLA)\n```\n\nExample:\n```text\npython build/build.py build --wheels=jaxlib --python_version=3.12 --verbose\n```\n\nExample:\n```text\npython build/build.py build --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt\n```\n\nExample:\n```text\npython build/build.py build --wheels=jax-cuda-plugin --cuda_version=12.3.2 \\\n--cudnn_version=9.1.1 --nccl_version=2.28.9\n```\n\nExample:\n```text\npython build/build.py build --wheels=jax-cuda-pjrt --cuda_version=12.3.2 \\\n--cudnn_version=9.1.1  --nccl_version=2.28.9\n```\n\nExample:\n```text\npython build/build.py build --wheels=jax-cuda-plugin \\\n--bazel_options=--repo_env=LOCAL_CUDA_PATH=\"/foo/bar/nvidia/cuda\" \\\n--bazel_options=--repo_env=LOCAL_CUDNN_PATH=\"/foo/bar/nvidia/cudnn\" \\\n--bazel_options=--repo_env=LOCAL_NCCL_PATH=\"/foo/bar/nvidia/nccl\"\n```\n\nExample:\n```text\npython build/build.py build --wheels=jaxlib --local_xla_path=/path/to/xla\n```\n\nExample:\n```text\npacman -S patch coreutils\n```\n\nExample:\n```text\npython .\\build\\build.py build --wheels=jaxlib\n```\n\nExample:\n```text\npython build/build.py build --python_version=3.12\n```\n\nExample:\n```text\n# Either add an entry to your `.bazelrc` file\nbuild --repo_env=HERMETIC_PYTHON_VERSION=3.12\n\n# OR pass it directly to your specific build command\nbazel build <target> --repo_env=HERMETIC_PYTHON_VERSION=3.12\n\n# OR set the environment variable globally in your shell:\nexport HERMETIC_PYTHON_VERSION=3.12\n```\n\nExample:\n```text\npython build/build.py requirements_update --python_version=3.12\n```\n\nExample:\n```text\nbazel run //build:requirements.update --repo_env=HERMETIC_PYTHON_VERSION=3.12\n```\n\nExample:\n```text\nbazel run //build:requirements.update --repo_env=HERMETIC_PYTHON_VERSION=3.12 -- --pre\n```\n\nExample:\n```text\necho -e \"\\n$(realpath jaxlib-0.4.27.dev20240416-cp312-cp312-manylinux_2_27_x86_64.whl)\" >> build/requirements.in\npython build/build.py requirements_update --python_version=3.12\n```\n\nExample:\n```text\npython build/build.py requirements_update --python_version=3.12 --nightly_update\n```\n\nExample:\n```text\nbazel run //build:requirements_nightly.update --repo_env=HERMETIC_PYTHON_VERSION=3.12\n```\n\nExample:\n```text\n./configure --prefix python\nmake -j12\nmake altinstall\ntar -czpf my_python.tgz python\n```\n\nExample:\n```text\n--repo_env=HERMETIC_PYTHON_URL=\"file:///local/path/to/my_python.tgz\"\n--repo_env=HERMETIC_PYTHON_SHA256=<file's_sha256_sum>\n\n# OR\n--repo_env=HERMETIC_PYTHON_URL=\"https://remote/url/to/my_python.tgz\"\n--repo_env=HERMETIC_PYTHON_SHA256=<file's_sha256_sum>\n\n# We assume that top-level folder in the tarball is called \"python\", if it is\n# something different just pass additional HERMETIC_PYTHON_PREFIX parameter\n--repo_env=HERMETIC_PYTHON_URL=\"https://remote/url/to/my_python.tgz\"\n--repo_env=HERMETIC_PYTHON_SHA256=<file's_sha256_sum>\n--repo_env=HERMETIC_PYTHON_PREFIX=\"my_python/install\"\n```\n\nExample:\n```text\n--repo_env=HERMETIC_REQUIREMENTS_LOCK=\"/absolute/path/to/custom_requirements_lock.txt\"\n```\n\nExample:\n```text\nbazel build <target>\n  --repo_env=HERMETIC_PYTHON_VERSION=3.13\n  --repo_env=HERMETIC_PYTHON_URL=\"https://github.com/indygreg/python-build-standalone/releases/download/20241016/cpython-3.13.0+20241016-x86_64-unknown-linux-gnu-install_only.tar.gz\"\n  --repo_env=HERMETIC_PYTHON_SHA256=\"2c8cb15c6a2caadaa98af51df6fe78a8155b8471cb3dd7b9836038e0d3657fb4\"\n```\n\nExample:\n```text\nbazel test <target>\n  --repo_env=HERMETIC_PYTHON_VERSION=3.13\n  --repo_env=HERMETIC_PYTHON_URL=\"file:///path/to/cpython.tar.gz\"\n  --repo_env=HERMETIC_PYTHON_PREFIX=\"prefix/to/strip/in/cython/tar/gz/archive\"\n  --repo_env=HERMETIC_PYTHON_SHA256=<sha256_sum>\n  --repo_env=HERMETIC_REQUIREMENTS_LOCK=\"/absolute/path/to/build:custom_requirements_lock.txt\"\n```\n\nExample:\n```text\nbazel test <target>\n  --repo_env=HERMETIC_PYTHON_VERSION=3.13\n  --repo_env=HERMETIC_REQUIREMENTS_LOCK=\"/absolute/path/to/build:custom_requirements_lock.txt\"\n```\n\nExample:\n```text\nrequirements = {\n  \"3.12\": \"//build:requirements_lock_3_12.txt\",\n  \"3.13\": \"//build:requirements_lock_3_13.txt\",\n  \"3.13-scenario1\": \"//build:scenario1_requirements_lock_3_13.txt\",\n  \"3.13-scenario2\": \"//build:scenario2_requirements_lock_3_13.txt\",\n},\n```\n\nExample:\n```text\n# To build with scenario1 dependencies:\nbazel test <target> --repo_env=HERMETIC_PYTHON_VERSION=3.13-scenario1\n\n# To build with scenario2 dependencies:\nbazel test <target> --repo_env=HERMETIC_PYTHON_VERSION=3.13-scenario2\n\n# To build with default dependencies:\nbazel test <target> --repo_env=HERMETIC_PYTHON_VERSION=3.13\n\n# To build with scenario1 dependencies and custom Python 3.13 interpreter:\nbazel test <target>\n  --repo_env=HERMETIC_PYTHON_VERSION=3.13-scenario1\n  --repo_env=HERMETIC_PYTHON_URL=\"file:///path/to/cpython.tar.gz\"\n  --repo_env=HERMETIC_PYTHON_SHA256=<sha256_sum>\n```\n\nExample:\n```text\npip install -e .  # installs jax\n```\n\nExample:\n```text\npython build/build.py build --wheels=jaxlib --configure_only\npython build/build.py build --wheels=jax-cuda-plugin --configure_only\npython build/build.py build --wheels=jax-rocm-plugin --configure_only\n```\n\nExample:\n```text\nbazel test //tests:cpu_tests //tests:backend_independent_tests\n```\n\nExample:\n```text\npython build/build.py build --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt --configure_only\n```\n\nExample:\n```text\necho -e \"\\njaxlib >= 0.4.26\" >> build/requirements.in\npython build/build.py requirements_update\n```\n\nExample:\n```text\necho -e \"\\n$(realpath jaxlib-0.4.26-cp312-cp312-manylinux_2_27_x86_64.whl)\" >> build/requirements.in\npython build/build.py requirements_update --python_version=3.12\n```\n\nExample:\n```text\nbazel test --//jax:build_jaxlib=false //tests:cpu_tests //tests:backend_independent_tests\n```\n\nExample:\n```text\nbazel test //tests:gpu_tests --local_test_jobs=4 --test_tag_filters=multiaccelerator --//jax:build_jaxlib=false --test_env=XLA_PYTHON_CLIENT_ALLOCATOR=platform\n```\n\nExample:\n```text\nNB_GPUS=2\nJOBS_PER_ACC=4\nJ=$((NB_GPUS * JOBS_PER_ACC))\nMULTI_GPU=\"--run_under $PWD/build/parallel_accelerator_execute.sh --test_env=JAX_ACCELERATOR_COUNT=${NB_GPUS} --test_env=JAX_TESTS_PER_ACCELERATOR=${JOBS_PER_ACC} --local_test_jobs=$J\"\nbazel test //tests:gpu_tests //tests:backend_independent_tests --test_env=XLA_PYTHON_CLIENT_PREALLOCATE=false --test_tag_filters=-multiaccelerator $MULTI_GPU\n```\n\nExample:\n```text\npytest -n auto tests\n```\n\nExample:\n```text\n# Bazel\nbazel test //tests/... --test_env=JAX_NUM_GENERATED_CASES=25`\n```\n\nExample:\n```text\n# pytest\nJAX_NUM_GENERATED_CASES=25 pytest -n auto tests\n```\n\nExample:\n```text\nJAX_ENABLE_X64=1 JAX_NUM_GENERATED_CASES=25 pytest -n auto tests\n```\n\nExample:\n```text\nJAX_NUM_GENERATED_CASES=5 python tests/lax_numpy_test.py\n```\n\nExample:\n```text\npython tests/lax_numpy_test.py --test_targets=\"testPad\"\n```\n\nExample:\n```text\nYou can reproduce this example by temporarily adding @reproduce_failure('6.97.4', b'AXicY2DAAAAAEwAB') as a decorator on your test case\n```\n\nExample:\n```text\nJAX_TRACEBACK_FILTERING=off XLA_FLAGS=--xla_force_host_platform_device_count=8 pytest -n auto --tb=short --doctest-glob='*.md' --doctest-glob='*.rst' docs --doctest-continue-on-failure --ignore=docs/multi_process.md\n```\n\nExample:\n```text\nJAX_TRACEBACK_FILTERING=off XLA_FLAGS=--xla_force_host_platform_device_count=8 pytest --doctest-modules jax/_src/numpy/lax_numpy.py\n```\n\nExample:\n```text\npip install pre-commit\npre-commit run pyrefly-check --all-files\n```\n\nExample:\n```text\npip install pre-commit\npre-commit run ruff --all-files\n```\n\nExample:\n```text\npip install -r docs/requirements.txt\n```\n\nExample:\n```text\nsphinx-build -b html docs docs/build/html -j auto\n```\n\nExample:\n```text\nsphinx-build -b html -D nb_execution_mode=off docs docs/build/html -j auto\n```\n\nExample:\n```text\npip install jupytext==1.16.4\njupytext --sync docs/notebooks/thinking_in_jax.ipynb\n```\n\nExample:\n```text\npip install pre-commit\npre-commit run jupytext --all-files\n```\n\nExample:\n```text\njupytext --set-formats ipynb,md:myst path/to/the/notebook.ipynb\n```\n\nExample:\n```text\nmkvirtualenv jax-docs  # A new virtualenv\nmkdir jax-docs  # A new directory\ncd jax-docs\ngit clone --no-single-branch --depth 50 https://github.com/jax-ml/jax\ncd jax\ngit checkout --force origin/test-docs\ngit clean -d -f -f\nworkon jax-docs\n\npython -m pip install --upgrade --no-cache-dir pip\npython -m pip install --upgrade --no-cache-dir -I Pygments==2.3.1 setuptools==41.0.1 docutils==0.14 mock==1.0.1 pillow==5.4.1 alabaster>=0.7,<0.8,!=0.7.5 commonmark==0.8.1 recommonmark==0.5.0 'sphinx<2' 'sphinx-rtd-theme<0.5' 'readthedocs-sphinx-ext<1.1'\npython -m pip install --exists-action=w --no-cache-dir -r docs/requirements.txt\ncd docs\npython `which sphinx-build` -T -E -b html -d _build/doctrees-readthedocs -D language=en . _build/html\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.759Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":54,"totalLines":355,"estimatedTokens":2219}}68{"id":"doc-pallas_async_operations_jax_documentation-de89ee7e","source":"documentation","title":"Pallas Async Operations — JAX documentation","url":"https://docs.jax.dev/en/latest/pallas/design/async_note.html","text":"Example:\n```text\ndef f(x):\n  y = ppermute(x)\n  z = x + 1\n  return y, z\n```\n\nExample:\n```text\ndef f(x):\n  fut = ppermute_start(x)\n  z = x + 1  # happens at the same time as ppermute\n  y = ppermute_done(fut)\n  return y, z\n```\n\nExample:\n```text\ndef ppermute_kernel(x_ref, y_ref, send_sem, recv_sem):\n  right_neighbor = ...\n  descriptor = pltpu.make_async_remote_copy(x_ref, y_ref, send_sem, recv_sem, device_id=right_neighbor)\n  descriptor.start()\n  descriptor.wait_send()\n  descriptor.wait_recv()\n\ndef ppermute(x):\n  return pl.pallas_call(ppermute_kernel, out_shape=x, ...)(x)\n```\n\nExample:\n```text\ndef add_one(x_ref, z_ref):\n  z_ref[...] = x_ref[...] + 1\n\ndef ppermute_add_one_kernel(x_ref, y_ref, z_ref, send_sem, recv_sem):\n  right_neighbor = ...\n  descriptor = pltpu.make_async_remote_copy(x_ref, y_ref, send_sem, recv_sem, device_id=right_neighbor)\n  descriptor.start()\n\n  # Explicitly schedule inner kernel between start/wait\n  pltpu.emit_pipeline(add_one)(x_ref, z_ref)\n\n  descriptor.wait_send()\n  descriptor.wait_recv()\n\ndef ppermute_and_add_one(x):\n  return pl.pallas_call(ppermute_add_one_kernel, out_shape=(x, x), ...)(x)\n```\n\nExample:\n```text\ndef ppermute_start_kernel(\n    in_ref, send_sem, recv_sem, out_ref, *, axis_name,\n):\n  axis_size = jax.lax.psum(1, axis_name)\n  left_neighbor = jax.lax.rem(\n      jax.lax.axis_index(axis_name) - 1 + axis_size, axis_size\n  )\n  right_neighbor = jax.lax.rem(jax.lax.axis_index(axis_name) + 1, axis_size)\n  barrier_sem = pltpu.get_barrier_semaphore()\n  pltpu.semaphore_signal(barrier_sem, device_id=left_neighbor)\n  pltpu.semaphore_wait(barrier_sem, 1)\n  pltpu.make_async_remote_copy(\n      in_ref, out_ref, send_sem, recv_sem, device_id=right_neighbor\n  ).start()\n\ndef ppermute_start(x, *, axis_name) -> tuple[Semaphore, Semaphore, Array]:\n  send_sem, recv_sem, out = pl.pallas_call(\n      functools.partial(ppermute_start_kernel, axis_name=axis_name),\n      out_shape=(\n          pltpu.SemaphoreType.DMA(()),\n          pltpu.SemaphoreType.DMA(()),\n          jax.ShapeDtypeStruct(\n              x.shape,\n              dtype=x.dtype,\n          ),\n      ),\n      in_specs=[\n          pl.BlockSpec(memory_space=pl.ANY),\n      ],\n      out_specs=(\n          pl.BlockSpec(memory_space=pltpu.SEMAPHORE),\n          pl.BlockSpec(memory_space=pltpu.SEMAPHORE),\n          pl.BlockSpec(memory_space=pl.ANY),\n      ),\n  )(x)\n  return send_sem, recv_sem, out\n```\n\nExample:\n```text\ndef ppermute_done_kernel(ref, send_sem, recv_sem, _):\n  pltpu.make_async_copy(ref, ref, send_sem).wait()\n  pltpu.make_async_copy(ref, ref, recv_sem).wait()\n\ndef ppermute_done(send_sem, recv_sem, out) ->Array:\n  out = pl.pallas_call(\n      ppermute_done_kernel,\n      out_shape=(\n          jax.ShapeDtypeStruct(\n              out.shape,\n              dtype=out.dtype,\n          ),\n      ),\n      in_specs=[\n          pl.BlockSpec(memory_space=pl.ANY),\n          pl.BlockSpec(memory_space=pltpu.SEMAPHORE),\n          pl.BlockSpec(memory_space=pltpu.SEMAPHORE),\n      ],\n      out_specs=pl.BlockSpec(memory_space=pl.ANY),\n      input_output_aliases={0:0}\n  )(out, send_sem, recv_sem)\n  return out\n```\n\nExample:\n```text\ndef f(x):\n  fut = ppermute_start(x)\n  z = x + 1\n  y = ppermute_done(fut)\n  return y, z\n```\n\nExample:\n```text\ndef f(x):\n  z = x + 1\n  fut = ppermute_start(x)\n  y = ppermute_done(fut)\n  return y, z\n\n# OR\n\ndef f(x):\n  fut = ppermute_start(x)\n  z = x + 1\n  y = ppermute_done(fut)\n  return y, z\n\n# OR\n\ndef f(x):\n  fut = ppermute_start(x)\n  y = ppermute_done(fut)\n  z = x + 1\n  return y, z\n```\n\nExample:\n```text\ndef f(x):\n  fut = ppermute_start(x)\n  x, fut = optimization_barrier((x, fut))  # x now depends on fut\n  z = x + 1\n  z, fut = optimization_barrier((z, fut)) # fut now depends on z\n  y = ppermute_done(fut)\n  return y, z\n```\n\nExample:\n```text\ndef f(x):\n  fut = ppermute_start(x)\n  z = x + 1\n  # XLA can free x here!\n  y = ppermute_done(fut)\n  return y, z\n```\n\nExample:\n```text\ndef ppermute_start_kernel(\n    in_ref, send_sem, recv_sem, out_ref, _, *, axis_name,\n):\n  axis_size = jax.lax.psum(1, axis_name)\n  left_neighbor = jax.lax.rem(\n      jax.lax.axis_index(axis_name) - 1 + axis_size, axis_size\n  )\n  right_neighbor = jax.lax.rem(jax.lax.axis_index(axis_name) + 1, axis_size)\n  barrier_sem = pltpu.get_barrier_semaphore()\n  pltpu.semaphore_signal(barrier_sem, device_id=left_neighbor)\n  pltpu.semaphore_wait(barrier_sem, 1)\n  pltpu.make_async_remote_copy(\n      in_ref, out_ref, send_sem, recv_sem, device_id=right_neighbor\n  ).start()\n\ndef ppermute_start(x, *, axis_name) -> tuple[Semaphore, Semaphore, Array, Array]:\n  send_sem, recv_sem, x, out = pl.pallas_call(\n      functools.partial(ppermute_start_kernel, axis_name=axis_name),\n      out_shape=(\n          pltpu.SemaphoreType.DMA(()),\n          pltpu.SemaphoreType.DMA(()),\n          jax.ShapeDtypeStruct(\n              x.shape,\n              dtype=x.dtype,\n          ),\n\t   jax.ShapeDtypeStruct(\n              x.shape,\n              dtype=x.dtype,\n          ),\n      ),\n      in_specs=[\n          pl.BlockSpec(memory_space=pl.ANY),\n      ],\n      out_specs=(\n          pl.BlockSpec(memory_space=pltpu.SEMAPHORE),\n          pl.BlockSpec(memory_space=pltpu.SEMAPHORE),\n          pl.BlockSpec(memory_space=pl.ANY),\n          pl.BlockSpec(memory_space=pl.ANY),\n      ),\n      input_output_aliases={0:2}\n  )(x)\n  return send_sem, recv_sem, x, out\n```\n\nExample:\n```text\ndef ppermute_done_kernel(_, ref, send_sem, recv_sem, _):\n  pltpu.make_async_copy(ref, ref, send_sem).wait()\n  pltpu.make_async_copy(ref, ref, recv_sem).wait()\n\ndef ppermute_done(send_sem, recv_sem, x, out) ->Array:\n  out = pl.pallas_call(\n      ppermute_done_kernel,\n      out_shape=(\n          jax.ShapeDtypeStruct(\n              out.shape,\n              dtype=out.dtype,\n          ),\n      ),\n      in_specs=[\n          pl.BlockSpec(memory_space=pl.ANY),\n          pl.BlockSpec(memory_space=pl.ANY),\n          pl.BlockSpec(memory_space=pltpu.SEMAPHORE),\n          pl.BlockSpec(memory_space=pltpu.SEMAPHORE),\n      ],\n      out_specs=pl.BlockSpec(memory_space=pl.ANY),\n      input_output_aliases={1:0}\n  )(x, out, send_sem, recv_sem)\n  return out\n```\n\nExample:\n```text\ndef f(x):\n  *sems, x ,out = ppermute_start(x)\n  z = x + 1\n  y = ppermute_done(*sems, x, out)\n  return y, z\n```\n\nExample:\n```text\ndef f():\n  x = jnp.arange(...)\n  y = add_one_inplace(x)\u000b  return y\n```\n\nExample:\n```text\ndef f():\n  x = jnp.arange(...)\n  y = add_one_inplace(x)\n  return y, x * 2 # another x consumer!\n```\n\nExample:\n```text\ndef f(x):\n  x2 = copy(x)\n  y = add_one_inplace(x2)\n  return y, x * 2\n```\n\nExample:\n```text\ndef f(x):\n  *sems, x2, y = ppermute_start(x)\n  z = x + 1\n  y = ppermute_done((*sems, x2, y))\n  return y, z\n```\n\nExample:\n```text\ndef f(x):\n  x2 = copy(x)\n  *sems, x2, y = ppermute_start(x2)\n  z = x + 1\n  y = ppermute_done((*sems, x2, y))\n  return y, z\n```\n\nExample:\n```text\ndef f(x):\n  *sems, x2, y = ppermute_start(x)\n  z = x2 + 1\n  y = ppermute_done((*sems, x2, y))\n  return y, z\n```\n\nExample:\n```text\ndef f(x):\n  def body(i, x):\n    fut = ppermute_start(x)\n    y = ppermute_done(fut)\n    return y\n  return fori_loop(0, 8, body, x)\n```\n\nExample:\n```text\ndef f(x):\n  def body(i, x):\n    *sems, x, y = ppermute_start(x)\n    y = ppermute_done(*sems, x, y)\n    return y\n  return fori_loop(0, 8, body, x)\n```\n\nExample:\n```text\ndef f(x):\n  def body(i, x):\n    *sems, x, y = ppermute_start(x)\n    y = ppermute_done((*sems, x, y))\n    return y\n  return fori_loop(0, 8, body, x)\n```\n\nExample:\n```text\ndef f(x):\n  def body(i, x):\n    x = copy(x)\n    *sems, x, y = ppermute_start(x)\n    y = ppermute_done((*sems, x, y))\n    return y\n  return fori_loop(0, 8, body, x)\n```\n\nExample:\n```text\ndef f(x):\n  def body(i, x):\n    *sems, x, x2 = ppermute_start(x)\n    x2 = ppermute_done((*sems, x, x2))\n\n    *sems, x2, y = ppermute_start(x2)\n    y = ppermute_done((*sems, x2, y))\n    return y\n  return fori_loop(0, 4, body, x)\n```\n\nExample:\n```text\ndef f(x):\n  def body(i, x):\n    fut = ppermute_start(x)\n    y = ppermute_done(fut)\n    return y\n  return fori_loop(0, 8, body, x, unroll=2)\n```\n\nExample:\n```text\ndef f(x):\n  fut = ppermute_start(x)\n  def body(i, fut):\n    x = ppermute_done(fut)\n    fut = ppermute_start(x)\n    return fut\n  fut = fori_loop(0, 7, body, fut)\n  return ppermute_done(fut)\n```\n\nExample:\n```text\ndef f(x):\n  fut = ppermute_start(x)\n  def body(i, fut):\n    *sems, x, out = fut\n    x = ppermute_done((*sems, x, out))\n    (*sems, x, out) = ppermute_start(x)\n    return (*sems, x, out)\n  (*sems, x, out) = fori_loop(0, 7, body, x)\n  return ppermute_done((*sems, x, out))\n```\n\nExample:\n```text\ndef f(x):\n  fut = ppermute_start(x)\n  def body(i, fut):\n    *sems, x, out = fut\n    out = copy(out)\n    x = ppermute_done((*sems, x, out))\n    (*sems, x, out) = ppermute_start(x)\n    return (*sems, x, out)\n  (*sems, x, out) = fori_loop(0, 7, body, x)\n  return ppermute_done((*sems, x, out))\n```\n\nExample:\n```text\ndef f(x):\n  fut = ppermute_start(x)\n  def body(i, fut):\n    x = ppermute_done(fut)\n    fut = ppermute_start(x)\n    return fut\n  fut = fori_loop(0, 7, body, x, unroll=2)\n  return ppermute_done(fut)\n```\n\nExample:\n```text\ndef f(x):\n  out = jnp.zeros_like(x)\n  fut = (*sems, x, out) = ppermute_start(x)\n  out = out + x\n  def body(i, carry):\n    out, fut = carry\n    x = ppermute_done(fut)\n    fut = (*sems, x, out) = ppermute_start(x)\n    out = out + x\n    return out, fut\n  out, fut = fori_loop(0, 7, body, (out, fut), unroll=2)\n  return out, ppermute_done(fut)\n```\n\nExample:\n```text\ndef ppermute_start_stateful(x_ref, y_ref) -> tuple[Semaphore, Semaphore]:\n  ...\n\ndef ppermute_done_stateful(send_sem, recv_sem, x_ref, y_ref) -> None:\n  ...\n```\n\nExample:\n```text\ndef f(x):\n  x_ref = make_ref(x)\n  y_ref = make_ref(zeros_like(x))\n  fut = ppermute_start_stateful(x_ref, y_ref)\n  ppermute_done_stateful(*fut, x_ref, y_ref)\n  return y_ref[...]\n```\n\nExample:\n```text\ndef f(x):\n  x_ref = make_ref(x)\n  y_ref = make_ref(zeros_like(x))\n  fut = ppermute_start_stateful(x_ref, y_ref)\n  x_ref[...] += 1\n  ppermute_done_stateful(*fut, x_ref, y_ref)\n  return y_ref[...]\n```\n\nExample:\n```text\ndef f(x):\n  x_ref = make_ref(x)\n  y_ref = make_ref(zeros_like(x))\n  def body(i, _):\n    fut = ppermute_start_stateful(x_ref, y_ref)\n    ppermute_done_stateful(*fut, x_ref, y_ref)\n    # Now switch to y_ref -> x_ref\n    fut = ppermute_start_stateful(y_ref, x_ref)\n    ppermute_done_stateful(*fut, y_ref, x_ref)\n  fori_loop(0, 8 // 2, body, None)\n  return x_ref[...]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.761Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":34,"totalLines":481,"estimatedTokens":2607}}69{"id":"doc-errors_jax_documentation-3df41e29","source":"documentation","title":"Errors — JAX documentation","url":"https://docs.jax.dev/en/latest/errors.html","text":"Example:\n```text\n>>> from functools import partial\n>>> from jax import jit\n>>> import jax.numpy as jnp\n>>> @jit\n... def func(x, axis):\n...   return x.min(axis)\n```\n\nExample:\n```text\n>>> func(jnp.arange(4), 0)  \nTraceback (most recent call last):\n    ...\nConcretizationTypeError: Abstract tracer value encountered where concrete\nvalue is expected: axis argument to jnp.min().\n```\n\nExample:\n```text\n>>> @jit(static_argnums=1)\n... def func(x, axis):\n...   return x.min(axis)\n\n>>> func(jnp.arange(4), 0)\nArray(0, dtype=int32)\n```\n\nExample:\n```text\n>>> @jit\n... def func(x):\n...     return jnp.where(x < 0)\n\n>>> func(jnp.arange(4))  \nTraceback (most recent call last):\n    ...\nConcretizationTypeError: Abstract tracer value encountered where concrete value is expected:\nThe error arose in jnp.nonzero.\n```\n\nExample:\n```text\n>>> @jit\n... def func(x):\n...     indices = jnp.where(x > 1)\n...     return x[indices].sum()\n\n>>> func(jnp.arange(4))  \nTraceback (most recent call last):\n    ...\nConcretizationTypeError: Abstract tracer value encountered where concrete\nvalue is expected: The error arose in jnp.nonzero.\n```\n\nExample:\n```text\n>>> @jit\n... def func(x):\n...   return jnp.where(x > 1, x, 0).sum()\n\n>>> func(jnp.arange(4))\nArray(5, dtype=int32)\n```\n\nExample:\n```text\n>>> with jax.debug_key_reuse(True):  \n...   key = jax.random.key(0)\n...   value = jax.random.uniform(key)\n...   new_value = jax.random.uniform(key)\n...\n---------------------------------------------------------------------------\nKeyReuseError                             Traceback (most recent call last)\n...\nKeyReuseError: Previously-consumed key passed to jit-compiled function at index 0\n```\n\nExample:\n```text\n>>> import jax\n>>> import jax.numpy as jnp\n\n>>> @jax.jit\n... def positive_values(x):\n...   return x[x > 0]\n\n>>> positive_values(jnp.arange(-5, 5))  \nTraceback (most recent call last):\n    ...\nNonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[10])\n```\n\nExample:\n```text\n>>> @jax.jit\n... def sum_of_positive(x):\n...   return x[x > 0].sum()\n\n>>> sum_of_positive(jnp.arange(-5, 5))  \nTraceback (most recent call last):\n    ...\nNonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[10])\n```\n\nExample:\n```text\n>>> @jax.jit\n... def sum_of_positive(x):\n...   return jnp.where(x > 0, x, 0).sum()\n\n>>> sum_of_positive(jnp.arange(-5, 5))\nArray(10, dtype=int32)\n```\n\nExample:\n```text\n>>> @jax.jit\n... def manual_clip(x):\n...   return x.at[x < 0].set(0)\n\n>>> manual_clip(jnp.arange(-2, 2))  \nTraceback (most recent call last):\n    ...\nNonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[4])\n```\n\nExample:\n```text\n>>> @jax.jit\n... def manual_clip(x):\n...   return jnp.where(x < 0, 0, x)\n\n>>> manual_clip(jnp.arange(-2, 2))\nArray([0, 0, 0, 1], dtype=int32)\n```\n\nExample:\n```text\n>>> from jax import jit\n>>> import numpy as np\n\n>>> @jit\n... def func(x):\n...   return np.sin(x)\n\n>>> func(np.arange(4))  \nTraceback (most recent call last):\n    ...\nTracerArrayConversionError: The numpy.ndarray conversion method\n__array__() was called on traced array with shape int32[4]\n```\n\nExample:\n```text\n>>> import jax.numpy as jnp\n>>> @jit\n... def func(x):\n...   return jnp.sin(x)\n\n>>> func(jnp.arange(4))\nArray([0.        , 0.84147096, 0.9092974 , 0.14112   ], dtype=float32)\n```\n\nExample:\n```text\n>>> x = np.arange(10)\n\n>>> @jit\n... def func(i):\n...   return x[i]\n\n>>> func(0)  \nTraceback (most recent call last):\n    ...\nTracerArrayConversionError: The numpy.ndarray conversion method\n__array__() was called on traced array with shape int32[0]\n```\n\nExample:\n```text\n>>> @jit\n... def func(i):\n...   return jnp.asarray(x)[i]\n\n>>> func(0)\nArray(0, dtype=int32)\n```\n\nExample:\n```text\n>>> from functools import partial\n>>> @jit(static_argnums=(0,))\n... def func(i):\n...   return x[i]\n\n>>> func(0)\nArray(0, dtype=int32)\n```\n\nExample:\n```text\n>>> from jax import jit\n>>> import jax.numpy as jnp\n>>> @jit\n... def func(x, y):\n...   return x if x.sum() < y.sum() else y\n\n>>> func(jnp.ones(4), jnp.zeros(4))  \nTraceback (most recent call last):\n    ...\nTracerBoolConversionError: Attempted boolean conversion of JAX Tracer [...]\n```\n\nExample:\n```text\n>>> @jit\n... def func(x, y):\n...   return jnp.where(x.sum() < y.sum(), x, y)\n\n>>> func(jnp.ones(4), jnp.zeros(4))\nArray([0., 0., 0., 0.], dtype=float32)\n```\n\nExample:\n```text\n>>> @jit\n... def func(x, normalize=True):\n...   if normalize:\n...     return x / x.sum()\n...   return x\n\n>>> func(jnp.arange(5), True)  \nTraceback (most recent call last):\n    ...\nTracerBoolConversionError: Attempted boolean conversion of JAX Tracer ...\n```\n\nExample:\n```text\n>>> from functools import partial\n>>> @jit(static_argnames=['normalize'])\n... def func(x, normalize=True):\n...   if normalize:\n...     return x / x.sum()\n...   return x\n\n>>> func(jnp.arange(5), True)\nArray([0. , 0.1, 0.2, 0.3, 0.4], dtype=float32)\n```\n\nExample:\n```text\n>>> @jit\n... def func(x):\n...   return min(x, 0)\n```\n\nExample:\n```text\n>>> func(2)  \nTraceback (most recent call last):\n    ...\nTracerBoolConversionError: Attempted boolean conversion of JAX Tracer ...\n```\n\nExample:\n```text\n>>> @jit\n... def func(x):\n...   return jnp.minimum(x, 0)\n```\n\nExample:\n```text\n>>> print(func(2))\n0\n```\n\nExample:\n```text\n>>> from jax import jit\n>>> import numpy as np\n\n>>> @jit\n... def func(x, axis):\n...   return np.split(x, 2, axis)\n\n>>> func(np.arange(4), 0)  \nTraceback (most recent call last):\n    ...\nTracerIntegerConversionError: The __index__() method was called on\ntraced array with shape int32[0]\n```\n\nExample:\n```text\n>>> from functools import partial\n>>> @jit(static_argnums=1)\n... def func(x, axis):\n...   return np.split(x, 2, axis)\n\n>>> func(np.arange(10), 0)\n[Array([0, 1, 2, 3, 4], dtype=int32),\n Array([5, 6, 7, 8, 9], dtype=int32)]\n```\n\nExample:\n```text\n>>> jit(lambda arr: np.split(arr, 2, 0))(np.arange(4))\n[Array([0, 1], dtype=int32), Array([2, 3], dtype=int32)]\n```\n\nExample:\n```text\n>>> import jax.numpy as jnp\n>>> from jax import jit\n\n>>> L = [1, 2, 3]\n\n>>> @jit\n... def func(i):\n...   return L[i]\n\n>>> func(0)  \nTraceback (most recent call last):\n    ...\nTracerIntegerConversionError: The __index__() method was called on\ntraced array with shape int32[0]\n```\n\nExample:\n```text\n>>> @jit\n... def func(i):\n...   return jnp.array(L)[i]\n\n>>> func(0)\nArray(1, dtype=int32)\n```\n\nExample:\n```text\n>>> from functools import partial\n>>> @jit(static_argnums=0)\n... def func(i):\n...   return L[i]\n\n>>> func(0)\nArray(1, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\n>>> from jax import jit\n>>> import jax.numpy as jnp\n\n>>> outs = []\n>>> @jit                   # 1\n... def side_effecting(x):\n...   y = x + 1            # 3\n...   outs.append(y)       # 4\n\n>>> x = 1\n>>> side_effecting(x)      # 2\n>>> outs[0] + 1            # 5  \nTraceback (most recent call last):\n    ...\nUnexpectedTracerError: Encountered an unexpected tracer.\n```\n\nExample:\n```text\n>>> from jax import jit\n>>> import jax.numpy as jnp\n\n>>> outs = []\n>>> @jit\n... def not_side_effecting(x):\n...   y = x+1\n...   return y\n\n>>> x = 1\n>>> y = not_side_effecting(x)\n>>> outs.append(y)\n>>> outs[0] + 1  # all good! no longer a leaked value.\nArray(3, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\n>>> from jax import jit\n>>> import jax.numpy as jnp\n\n>>> outs = []\n>>> @jit\n... def side_effecting(x):\n...   y = x+1\n...   outs.append(y)\n\n>>> x = 1\n>>> with jax.checking_leaks():\n...   y = side_effecting(x)  \nTraceback (most recent call last):\n    ...\nException: Leaked Trace\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.762Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":34,"totalLines":410,"estimatedTokens":1880}}70{"id":"doc-introduction_to_multi_controller_jax_aka_multi_p-27416b88","source":"documentation","title":"Introduction to multi-controller JAX (aka multi-process/multi-host JAX) — JAX documentation","url":"https://docs.jax.dev/en/latest/multi_process.html","text":"Example:\n```text\n# call this file toy.py, to be run in each process simultaneously\n\nimport jax\nimport jax.numpy as jnp\nfrom jax.sharding import NamedSharding, PartitionSpec as P\nimport numpy as np\n\n# in this example, get multi-process parameters from sys.argv\nimport sys\nproc_id = int(sys.argv[1])\nnum_procs = int(sys.argv[2])\n\n# initialize the distributed system\njax.distributed.initialize('localhost:10000', num_procs, proc_id)\n\n# this example assumes 8 devices total\nassert jax.device_count() == 8\n\n# make a 2D mesh that refers to devices from all processes\nmesh = jax.make_mesh((4, 2), ('i', 'j'))\n\n# create some toy data\nglobal_data = np.arange(32).reshape((4, 8))\n\n# make a process- and device-spanning array from our toy data\nsharding = NamedSharding(mesh, P('i', 'j'))\nglobal_array = jax.device_put(global_data, sharding)\nassert global_array.shape == global_data.shape\n\n# each process has different shards of the global array\nfor shard in global_array.addressable_shards:\n  print(f\"device {shard.device} has local data {shard.data}\")\n\n# apply a simple computation, automatically partitioned\nglobal_result = jnp.sum(jnp.sin(global_array))\nprint(f'process={proc_id} got result: {global_result}')\n```\n\nExample:\n```text\nexport JAX_NUM_CPU_DEVICES=2\nnum_processes=4\n\nrange=$(seq 0 $(($num_processes - 1)))\n\nfor i in $range; do\n  python toy.py $i $num_processes > /tmp/toy_$i.out &\ndone\n\nwait\n\nfor i in $range; do\n  echo \"=================== process $i output ===================\"\n  cat /tmp/toy_$i.out\n  echo\ndone\n```\n\nExample:\n```text\n=================== process 0 output ===================\ndevice TFRT_CPU_0 has local data [[0 1 2 3]]\ndevice TFRT_CPU_1 has local data [[4 5 6 7]]\nprocess=0 got result: -0.12398731708526611\n\n=================== process 1 output ===================\ndevice TFRT_CPU_131072 has local data [[ 8  9 10 11]]\ndevice TFRT_CPU_131073 has local data [[12 13 14 15]]\nprocess=1 got result: -0.12398731708526611\n\n=================== process 2 output ===================\ndevice TFRT_CPU_262144 has local data [[16 17 18 19]]\ndevice TFRT_CPU_262145 has local data [[20 21 22 23]]\nprocess=2 got result: -0.12398731708526611\n\n=================== process 3 output ===================\ndevice TFRT_CPU_393216 has local data [[24 25 26 27]]\ndevice TFRT_CPU_393217 has local data [[28 29 30 31]]\nprocess=3 got result: -0.12398731708526611\n```\n\nExample:\n```text\nJAX_NUM_CPU_DEVICES=8 python toy.py 0 1\n```\n\nExample:\n```text\ndevice TFRT_CPU_0 has local data [[0 1 2 3]]\ndevice TFRT_CPU_1 has local data [[4 5 6 7]]\ndevice TFRT_CPU_2 has local data [[ 8  9 10 11]]\ndevice TFRT_CPU_3 has local data [[12 13 14 15]]\ndevice TFRT_CPU_4 has local data [[16 17 18 19]]\ndevice TFRT_CPU_5 has local data [[20 21 22 23]]\ndevice TFRT_CPU_6 has local data [[24 25 26 27]]\ndevice TFRT_CPU_7 has local data [[28 29 30 31]]\nprocess=0 got result: -0.12398731708526611\n```\n\nExample:\n```text\n# In file gpu_example.py...\n\nimport jax\nimport sys\n\n# Get the coordinator_address, process_id, and num_processes from the command line.\ncoord_addr = sys.argv[1]\nproc_id = int(sys.argv[2])\nnum_procs = int(sys.argv[3])\n\n# Initialize the GPU machines.\njax.distributed.initialize(coordinator_address=coord_addr,\n                           num_processes=num_procs,\n                           process_id=proc_id)\nprint(\"process id =\", jax.process_index())\nprint(\"global devices =\", jax.devices())\nprint(\"local devices =\", jax.local_devices())\n```\n\nExample:\n```text\nprocess id = 0\nglobal devices = [CudaDevice(id=0), CudaDevice(id=1), CudaDevice(id=2), CudaDevice(id=3), CudaDevice(id=4), CudaDevice(id=5), CudaDevice(id=6), CudaDevice(id=7)]\nlocal devices = [CudaDevice(id=0), CudaDevice(id=1)]\n```\n\nExample:\n```text\nprocess id = 1\nglobal devices = [CudaDevice(id=0), CudaDevice(id=1), CudaDevice(id=2), CudaDevice(id=3), CudaDevice(id=4), CudaDevice(id=5), CudaDevice(id=6), CudaDevice(id=7)]\nlocal devices = [CudaDevice(id=2), CudaDevice(id=3)]\n```\n\nExample:\n```text\n$ TPU_NAME=jax-demo\n$ EXTERNAL_IPS=$(gcloud compute tpus tpu-vm describe $TPU_NAME --zone 'us-central1-a' \\\n                 | grep externalIp | cut -d: -f2)\n$ cat << EOF > demo.py\nimport jax\njax.distributed.initialize()\nif jax.process_index() == 0:\n  print(jax.devices())\nEOF\n$ echo $EXTERNAL_IPS | xargs -n 1 -P 0 bash -c '\nscp demo.py $0:\nssh $0 \"pip -q install -U jax[tpu]\"\nssh $0 \"python demo.py\" '\n```\n\nExample:\n```text\n[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=2, process_index=1, coords=(2,0,0), core_on_chip=0), TpuDevice(id=3, process_index=1, coords=(3,0,0), core_on_chip=0), TpuDevice(id=6, process_index=1, coords=(2,1,0), core_on_chip=0), TpuDevice(id=7, process_index=1, coords=(3,1,0), core_on_chip=0), TpuDevice(id=8, process_index=2, coords=(0,2,0), core_on_chip=0), TpuDevice(id=9, process_index=2, coords=(1,2,0), core_on_chip=0), TpuDevice(id=12, process_index=2, coords=(0,3,0), core_on_chip=0), TpuDevice(id=13, process_index=2, coords=(1,3,0), core_on_chip=0), TpuDevice(id=10, process_index=3, coords=(2,2,0), core_on_chip=0), TpuDevice(id=11, process_index=3, coords=(3,2,0), core_on_chip=0), TpuDevice(id=14, process_index=3, coords=(2,3,0), core_on_chip=0), TpuDevice(id=15, process_index=3, coords=(3,3,0), core_on_chip=0)]\n```\n\nExample:\n```text\napiVersion: jobset.x-k8s.io/v1alpha2\nkind: JobSet\nmetadata:\n  name: jaxjob\nspec:\n  replicatedJobs:\n  - name: workers\n    template:\n      spec:\n        parallelism: 2\n        completions: 2\n        backoffLimit: 0\n        template:\n          spec:\n            serviceAccountName: jax-job-sa  # kubectl apply -f svc-acct.yaml\n            restartPolicy: Never\n            imagePullSecrets:\n              # https://k8s.io/docs/tasks/configure-pod-container/pull-image-private-registry/\n            - name: null\n            containers:\n            - name: main\n              image: null  # e.g. ghcr.io/nvidia/jax:jax\n              imagePullPolicy: Always\n              resources:\n                limits:\n                  cpu: 1\n                  # https://k8s.io/docs/tasks/manage-gpus/scheduling-gpus/\n                  nvidia.com/gpu: null\n              command: \n                - python\n              args:\n                - -c\n                - |\n                  import jax\n                  jax.distributed.initialize()\n                  print(jax.devices())\n                  print(jax.local_devices())\n                  assert jax.process_count() > 1\n                  assert len(jax.devices()) > len(jax.local_devices())\n```\n\nExample:\n```text\n$ kubectl apply -f example.yaml\n$ kubectl get pods -l jobset.sigs.k8s.io/jobset-name=jaxjob\nNAME                       READY   STATUS      RESTARTS   AGE\njaxjob-workers-0-0-xpx8l   0/1     Completed   0          8m32s\njaxjob-workers-0-1-ddkq8   0/1     Completed   0          8m32s\n```\n\nExample:\n```text\n$ kubectl logs -l jobset.sigs.k8s.io/jobset-name=jaxjob\n[CudaDevice(id=0), CudaDevice(id=1)]\n[CudaDevice(id=0)]\n[CudaDevice(id=0), CudaDevice(id=1)]\n[CudaDevice(id=1)]\n```\n\nExample:\n```text\nfrom jax.sharding import Mesh\nmesh = Mesh(jax.devices(), ('a',))\n\n# in this case, the same as\nmesh = jax.make_mesh((jax.device_count(),), ('a',))  # use this in practice\n```\n\nExample:\n```text\nMesh(create_hybrid_device_mesh((1, devices_per_slice), (num_slices, 1)), axis_names=(\"dcn\", \"ici\"))\n```\n\nExample:\n```text\narr = jax.device_put(jnp.ones((32, 32)), NamedSharding(mesh, P('a')))\nif jax.process_index() == 0:\n  jax.debug.visualize_array_sharding(arr)\n```\n\nExample:\n```text\n┌───────────────────────┐\n│         TPU 0         │\n├───────────────────────┤\n│         TPU 1         │\n├───────────────────────┤\n│         TPU 4         │\n├───────────────────────┤\n│         TPU 5         │\n├───────────────────────┤\n│         TPU 2         │\n├───────────────────────┤\n│         TPU 3         │\n├───────────────────────┤\n│         TPU 6         │\n├───────────────────────┤\n│         TPU 7         │\n├───────────────────────┤\n│         TPU 8         │\n├───────────────────────┤\n│         TPU 9         │\n├───────────────────────┤\n│        TPU 12         │\n├───────────────────────┤\n│        TPU 13         │\n├───────────────────────┤\n│        TPU 10         │\n├───────────────────────┤\n│        TPU 11         │\n├───────────────────────┤\n│        TPU 14         │\n├───────────────────────┤\n│        TPU 15         │\n└───────────────────────┘\n```\n\nExample:\n```text\nmesh = jax.make_mesh((jax.device_count() // 2, 2), ('a', 'b'))\n\ndef device_put(x, spec):\n  return jax.device_put(x, NamedSharding(mesh, spec))\n\n# construct global arrays by sharding over the global mesh\nx = device_put(jnp.ones((4096, 2048)), P('a', 'b'))\ny = device_put(jnp.ones((2048, 4096)), P('b', None))\n\n# run a distributed matmul\nz = jax.nn.relu(x @ y)\n\n# inspect the sharding of the result\nif jax.process_index() == 0:\n  jax.debug.visualize_array_sharding(z)\n  print()\n  print(z.sharding)\n```\n\nExample:\n```text\n┌───────────────────────┐\n│        TPU 0,1        │\n├───────────────────────┤\n│        TPU 4,5        │\n├───────────────────────┤\n│        TPU 8,9        │\n├───────────────────────┤\n│       TPU 12,13       │\n├───────────────────────┤\n│        TPU 2,3        │\n├───────────────────────┤\n│        TPU 6,7        │\n├───────────────────────┤\n│       TPU 10,11       │\n├───────────────────────┤\n│       TPU 14,15       │\n└───────────────────────┘\n\nNamedSharding(mesh=Mesh('a': 8, 'b': 2), spec=PartitionSpec('a',), memory_kind=device)\n```\n\nExample:\n```text\nRuntimeError: Fetching value for `jax.Array` that spans non-addressable (non process local) devices is not possible. You can use `jax.experimental.multihost_utils.process_allgather` to print the global array or use `.addressable_shards` method of jax.Array to inspect the addressable (process local) shards.\n```\n\nExample:\n```text\nw = device_put(z, P(None, None))\nif jax.process_index() == 0:\n  print(w)\n```\n\nExample:\n```text\nnum_devices = jax.device_count() // 2\nmesh = jax.make_mesh((num_devices,), ('a',),\n                     devices=jax.devices()[num_devices:],\n                     axis_types=(jax.sharding.AxisType.Explicit,))\nsharding = NamedSharding(mesh, P('a'))\n\ndata = np.arange(64).reshape((8, 8))\nx = jax.device_put(data, sharding)\n\n# inspect the sharding of the result\nif jax.process_index() == 0:\n  jax.debug.visualize_array_sharding(x)\n  print()\n  print(x.sharding)\n\n# inspect the data local to each host\nprint(f\"Devices attached to process {jax.process_index()}: {jax.local_devices()}\")\nprint(f\"Addressable data for process {jax.process_index()}:\")\nfor shard in x.addressable_shards:\n  print(f\"device {shard.device} has local data {shard.data}\")\n```\n\nExample:\n```text\n┌───────────────────────┐\n│         TPU 8         │\n├───────────────────────┤\n│         TPU 9         │\n├───────────────────────┤\n│        TPU 10         │\n├───────────────────────┤\n│        TPU 11         │\n├───────────────────────┤\n│        TPU 15         │\n├───────────────────────┤\n│        TPU 14         │\n├───────────────────────┤\n│        TPU 13         │\n├───────────────────────┤\n│        TPU 12         │\n└───────────────────────┘\n\nNamedSharding(mesh=Mesh('a': 8, axis_types=(Explicit,)), spec=PartitionSpec('a',), memory_kind=device)\n```\n\nExample:\n```text\nDevices attached to process 0: [TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(1,1,0), core_on_chip=0)]\n\nAddressable data for process 0:\n```\n\nExample:\n```text\nDevices attached to process 3: [TpuDevice(id=10, process_index=3, coords=(2,2,0), core_on_chip=0), TpuDevice(id=11, process_index=3, coords=(3,2,0), core_on_chip=0), TpuDevice(id=14, process_index=3, coords=(2,3,0), core_on_chip=0), TpuDevice(id=15, process_index=3, coords=(3,3,0), core_on_chip=0)]\nAddressable data for process 3\ndevice TPU_10(process=3,(2,2,0,0)) has local data [[16 17 18 19 20 21 22 23]]\ndevice TPU_11(process=3,(3,2,0,0)) has local data [[24 25 26 27 28 29 30 31]]\ndevice TPU_15(process=3,(3,3,0,0)) has local data [[32 33 34 35 36 37 38 39]]\ndevice TPU_14(process=3,(2,3,0,0)) has local data [[40 41 42 43 44 45 46 47]]\n```\n\nExample:\n```text\nresult = jnp.sum(jnp.sin(x))\nprint(f\"process={jax.process_index()} got result: {result}\")\n```\n\nExample:\n```text\nprocess=2 got result: Array(0.09658563, dtype=float32)\n```\n\nExample:\n```text\nprocess=0 got result: Array(shape=(), dtype=float32)\n```\n\nExample:\n```text\n# Create a sharding that contains half of the global devices for the first\n# stage of the pipeline.\nnum_devices = jax.device_count() // 2\nmesh_first_half = jax.make_mesh((num_devices,), ('a',),\n                                devices=jax.devices()[:num_devices],\n                                axis_types=(jax.sharding.AxisType.Explicit,))\nsharding_first_half = NamedSharding(mesh_first_half, P('a'))\n\n# Create a sharding that contains the other half of the devices.\nmesh_second_half = jax.make_mesh((num_devices,), ('a',),\n                                 devices=jax.devices()[num_devices:],\n                                 axis_types=(jax.sharding.AxisType.Explicit,))\nsharding_second_half = NamedSharding(mesh_second_half, P('a'))\n\n# Place the input data on the first mesh.\ndata = np.arange(64).reshape((8, 8))\nx = jax.device_put(data, sharding_first_half)\n\n# `f` is the first stage of the pipeline.\n@jax.jit\ndef f(x):\n  # Arbitrary JAX computation.\n  return x\n\n# `g` is the second stage of the pipeline\n@jax.jit\ndef g(x):\n  # More JAX operations.\n  return x\n\n# Run the first stage on the first set of devices.\ny = f(x)\n\n# Transfer the data to the second set of devices.\n# `device_put` must be called in all processes that participate in either\n# `y.sharding` or `sharding_second_half`, so it's important to call `y = f(x)`\n# in all processes -- not just those that participate in the first stage -- so\n# that we always have a reference to `y`.\nz = jax.device_put(y, sharding_second_half)\n\n# Run the second stage on the second set of devices.\nresult = g(z)\n```\n\nExample:\n```text\n# Pipeline stage functions. Each stage will run on a different device.\npipeline_stages = [f, g, f, g]\ndevices = jax.devices()[:4]\n\nmicrobatches = [np.arange(512**2).reshape((512, 512)) for _ in range(12)]\n\n# Each microbatch is enqueued on each device sequentially, but each device\n# conceptually has an independent queue of computations and transfers which can\n# run in parallel across queues. For example, because there are no data\n# dependencies between the microbatches, device 0 will immediately start a new\n# microbatch once the previous is finished, overlapping with the `device_put` to\n# device 1.\nresults = []\nfor mb in microbatches:\n  for d, s in zip(devices, pipeline_stages):\n    mb = jax.device_put(mb, d)\n    mb = s(mb)\n  results.append(mb)\n```\n\nExample:\n```text\n# target (micro)batch size across the whole cluster\nbatch_size = 1024\n# how many examples each process should load per batch\nper_process_batch_size = batch_size // jax.process_count()\n# how many examples each device will process per batch\nper_device_batch_size = batch_size // jax.device_count()\n\n# make a data-parallel mesh and sharding\nmesh = jax.make_mesh((jax.device_count(),), ('batch'))\nsharding = NamedSharding(mesh, P('batch'))\n\n# our \"data loader\". each process loads a different set of \"examples\".\nprocess_batch = np.random.rand(per_process_batch_size, 2048, 42)\n\n# assemble a global array containing the per-process batches from all processes\nglobal_batch = jax.make_array_from_process_local_data(sharding, process_batch)\n\n# sanity check that everything got sharded correctly\nassert global_batch.shape[0] == batch_size\nassert process_batch.shape[0] == per_process_batch_size\nassert global_batch.addressable_shards[0].data.shape[0] == per_device_batch_size\n```\n\nExample:\n```text\nshape = (jax.process_count(), jax.local_device_count())\nmesh = jax.make_mesh(shape, ('i', 'j'))\nsharding = NamedSharding(mesh, P('i', 'j'))\n\n# manually create per-device data equivalent to np.arange(jax.device_count())\n# i.e. each device will get a single scalar value from 0..N\nlocal_arrays = [\n    jax.device_put(\n        jnp.array([[jax.process_index() * jax.local_device_count() + i]]),\n        device)\n    for i, device in enumerate(jax.local_devices())\n]\n\n# assemble a global array from the local_arrays across all processes\nglobal_array = jax.make_array_from_single_device_arrays(\n    shape=shape,\n    sharding=sharding,\n    arrays=local_arrays)\n\n# sanity check\nassert (np.all(\n    jax.experimental.multihost_utils.process_allgather(global_array) ==\n    np.arange(jax.device_count()).reshape(global_array.shape)))\n```\n\nExample:\n```text\nnum_participating_processes = 2\nshape = (num_participating_processes, jax.local_device_count())\ndevices = (jax.local_devices(process_index=0) +\n           jax.local_devices(process_index=1))\nmesh = jax.make_mesh(shape, ('i', 'j'),\n                     axis_types=(jax.sharding.AxisType.Explicit,) * 2,\n                     devices=devices)\nsharding = NamedSharding(mesh, P('i', 'j'))\n\n# manually create per-device data in processes 0 and 1.\nif jax.process_index() in (0, 1):\n  local_arrays = [\n      jax.device_put(\n          jnp.array([[jax.process_index() * jax.local_device_count() + i]]),\n          device)\n      for i, device in enumerate(jax.local_devices())\n  ]\nelse:\n  local_arrays = []\n\n# assemble an array from the local_arrays across processes 0 and 1\narray = jax.make_array_from_single_device_arrays(\n    shape=shape,\n    sharding=sharding,\n    arrays=local_arrays,\n    dtype=jnp.int32)\n\n# sanity check\nif jax.process_index() in (0, 1):\n  for shard in array.addressable_shards:\n    assert shard.data.size == 1\nelse:\n  assert not array.addressable_shards\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.765Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":33,"totalLines":571,"estimatedTokens":4486}}71{"id":"doc-jax_memories_and_host_offloading_jax_documentati-2ef65a5a","source":"documentation","title":"JAX Memories and Host Offloading — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/host-offloading.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.sharding import Mesh, NamedSharding, PartitionSpec as P\nimport numpy as np\n\n# Create mesh\n# 1x1 mesh represents a single device with two named dimensions (x and y)\nmesh = Mesh(np.array(jax.devices()[0]).reshape(1, 1), ('x', 'y'))\n\n# Device sharding - partitions data along x and y dimensions\ns_dev = NamedSharding(mesh, P('x', 'y'), memory_kind=\"device\")\n\n# Host sharding - same partitioning but in pinned host memory\ns_host = s_dev.with_memory_kind('pinned_host')\n\nprint(s_dev)   # Shows device memory sharding\nprint(s_host)  # Shows pinned host memory sharding\n```\n\nExample:\n```text\nNamedSharding(mesh=Mesh('x': 1, 'y': 1), spec=PartitionSpec('x', 'y'), memory_kind=device)\nNamedSharding(mesh=Mesh('x': 1, 'y': 1), spec=PartitionSpec('x', 'y'), memory_kind=pinned_host)\n```\n\nExample:\n```text\n# Create a 2x4 array\narr = jnp.arange(8.0).reshape(2, 4)\n\n# Move arrays to different memory locations based on sharding objects\narr_host = jax.device_put(arr, s_host)  # Places in pinned host memory\narr_dev = jax.device_put(arr, s_dev)    # Places in device memory\n\n# Verify memory locations\nprint(arr_host.sharding.memory_kind)  # Output: pinned_host\nprint(arr_dev.sharding.memory_kind)   # Output: device\n```\n\nExample:\n```text\npinned_host\ndevice\n```\n\nExample:\n```text\nf = jax.jit(lambda x:x, out_shardings=s_dev)\nout_dev = f(arr_host)\nprint(\"Result value of H2D: \\n\", out_dev)\n```\n\nExample:\n```text\nResult value of H2D: \n [[0. 1. 2. 3.]\n [4. 5. 6. 7.]]\n```\n\nExample:\n```text\n# Instead of the lambda function, add_func can be defined explicitly\n# move data to device before computation\ndef add_func(x):  # Move data to device and add one\n  x = jax.device_put(x, s_dev)\n  return x + 1\n\nf = jax.jit(add_func, out_shardings=s_dev)\nout_dev = f(arr_host)\nprint(\"Result value of H2D and add 1 in device memory: \\n\", out_dev)\n```\n\nExample:\n```text\nResult value of H2D and add 1 in device memory: \n [[1. 2. 3. 4.]\n [5. 6. 7. 8.]]\n```\n\nExample:\n```text\nf = jax.jit(lambda x: x, out_shardings=s_host)\nout_host = f(arr_dev)      # Input arrays in the device memory while output arrays in the host memory\nprint(\"Result value of D2H: \\n\", out_host)\n```\n\nExample:\n```text\nResult value of D2H: \n [[0. 1. 2. 3.]\n [4. 5. 6. 7.]]\n```\n\nExample:\n```text\n# Initialize input and weights with small values (0.0001)\ninput = jnp.ones((256, 256), dtype=jnp.float32) * 0.001  # Input matrix: 256 x 256\nw1 = jnp.ones((10, 256, 1024), dtype=jnp.float32) * 0.001 # 10 layers of 256 x 1024 matrices\nw2 = jnp.ones((10, 1024, 256), dtype=jnp.float32) * 0.001 # 10 layers of 1024 x 256 matrices\n\ndef two_layers(x, w):\n  # Simple two-layer linear transformation\n  w1, w2 = w\n  y = x @ w1\n  return y @ w2, None\n\ndef scanned(w, x):\n  # Applies the layer function 10 times using JAX's scan operation\n  # Input: w (tuple of weight matrices), x (input matrix)\n  # Output: sum of the final layer's output\n  result = jax.lax.scan(two_layers, x, w)[0]\n  return jnp.sum(result)\n\n# Compile and compute gradients of the scanned function\nf = jax.jit(jax.grad(scanned))  # Apply JIT compilation to gradient computation\n\n# Analyze memory usage\ncompiled_step = f.lower((w1, w2), input).compile()\ncompiled_stats = compiled_step.memory_analysis()\n\nif compiled_stats is not None:\n  # Calculate total memory usage including temporary storage, arguments, and outputs\n  # Subtract alias size to avoid double-counting memory shared between different components\n  total = compiled_stats.temp_size_in_bytes + compiled_stats.argument_size_in_bytes \\\n      + compiled_stats.output_size_in_bytes - compiled_stats.alias_size_in_bytes\n  print(f\"Temp size: {compiled_stats.temp_size_in_bytes / (1024**2):.2f} MB\")\n  print(f\"Argument size: {compiled_stats.argument_size_in_bytes / (1024**2):.2f} MB\")\n  print(f\"Total size: {total/(1024**2):.2f} MB\")\n\n# Execute the function and print sample results\nresult = f((w1, w2), input)     # Execute the function with weights and input\nprint(\"Sample of results: \", result[0][0, 0, :5])\n```\n\nExample:\n```text\nTemp size: 17.25 MB\nArgument size: 20.25 MB\nTotal size: 57.50 MB\nSample of results:  [3.8312336e-07 3.8312336e-07 3.8312336e-07 3.8312336e-07 3.8312336e-07]\n```\n\nExample:\n```text\nfrom jax.ad_checkpoint import checkpoint_name\n\ndef layer_name(x, w):\n  w1, w2 = w\n  x = checkpoint_name(x, \"x\")\n  y = x @ w1\n  return y @ w2, None\n```\n\nExample:\n```text\nfrom jax import checkpoint_policies as cp\n\npolicy = cp.save_and_offload_only_these_names(\n    names_which_can_be_saved=[],          # No values stored on device\n    names_which_can_be_offloaded=[\"x\"],   # Offload activations labeled \"x\"\n    offload_src=\"device\",                 # Move from device memory\n    offload_dst=\"pinned_host\"             # To pinned host memory\n)\n```\n\nExample:\n```text\ndef scanned(w, x):\n  remat_layer = jax.remat(layer_name,\n                          policy=policy,     # Use our offloading policy\n                          prevent_cse=False) # Allow CSE optimizations\n  result = jax.lax.scan(remat_layer, x, w)[0]\n  return jnp.sum(result)\n\n# Initialize input and weights with small values (0.0001)\ninput = jnp.ones((256, 256), dtype=jnp.float32) * 0.001  # Input matrix: 256 x 256\nw1 = jnp.ones((10, 256, 1024), dtype=jnp.float32) * 0.001 # 10 layers of 256 x 1024 matrices\nw2 = jnp.ones((10, 1024, 256), dtype=jnp.float32) * 0.001 # 10 layers of 1024 x 256 matrices\n\n# Compile and compute gradients of the scanned function\nf = jax.jit(jax.grad(scanned))  # Apply JIT compilation to gradient computation\n\n# Analyze memory usage\ncompiled_step = f.lower((w1, w2), input).compile()\ncompiled_stats = compiled_step.memory_analysis()\n\nif compiled_stats is not None:\n  total = compiled_stats.temp_size_in_bytes + compiled_stats.argument_size_in_bytes \\\n      + compiled_stats.output_size_in_bytes - compiled_stats.alias_size_in_bytes\n  print(f\"Temp size: {compiled_stats.temp_size_in_bytes / (1024**2):.2f} MB\")\n  print(f\"Argument size: {compiled_stats.argument_size_in_bytes / (1024**2):.2f} MB\")\n  print(f\"Total size: {total/(1024**2):.2f} MB\")\n\nresult_activation = f((w1, w2), input)     # Execute the function with weights and input\n# Verify numerical correctness\nare_close = jnp.allclose(\n    result_activation[0],    # Result from activation offloading only\n    result[0],         # Result from both activation and parameter offloading\n    rtol=1e-5,\n    atol=1e-5\n)\nprint(f\"Results match within tolerance: {are_close}\")\nprint(\"Sample of results: \", result_activation[0][0, 0, :5])\n```\n\nExample:\n```text\nTemp size: 6.50 MB\nArgument size: 20.25 MB\nTotal size: 46.75 MB\nResults match within tolerance: True\nSample of results:  [3.8312336e-07 3.8312336e-07 3.8312336e-07 3.8312336e-07 3.8312336e-07]\n```\n\nExample:\n```text\n# Hybrid version: Both activation and parameter offloading\ndef hybrid_layer(x, w):\n  # Move model parameters w1 and w2 to device memory via device_put\n  w1, w2 = jax.tree.map(lambda x: jax.device_put(x, s_dev), w)\n  x = checkpoint_name(x, \"x\")  # Offload activation x to host memory\n  y = x @ w1\n  return y @ w2, None\n\ndef hybrid_scanned(w, x):\n  remat_layer = jax.remat(hybrid_layer,     # Use hybrid_layer instead of layer\n                          policy=policy,     # Use offloading policy\n                          prevent_cse=False) # Allow CSE optimizations\n  result = jax.lax.scan(remat_layer, x, w)[0]\n  return jnp.sum(result)\n\n# Move model parameters w1 and w2 to the host via device_put\n# Initialize input and weights with small values (0.0001)\nwh1 = jax.device_put(w1, s_host)\nwh2 = jax.device_put(w2, s_host)\n\n# Compile and compute gradients of the scanned function\nf = jax.jit(jax.grad(hybrid_scanned))  # Apply JIT compilation to gradient computation\n\n# Analyze memory usage\ncompiled_step = f.lower((wh1, wh2), input).compile()\ncompiled_stats = compiled_step.memory_analysis()\n\nif compiled_stats is not None:\n  total = compiled_stats.temp_size_in_bytes + compiled_stats.argument_size_in_bytes \\\n      + compiled_stats.output_size_in_bytes - compiled_stats.alias_size_in_bytes\n  print(f\"Temp size: {compiled_stats.temp_size_in_bytes / (1024**2):.2f} MB\")\n  print(f\"Argument size: {compiled_stats.argument_size_in_bytes / (1024**2):.2f} MB\")\n  print(f\"Total size: {total / (1024**2):.2f} MB\")\n\nresult_both = f((wh1, wh2), input) # Execute with both activation and parameter offloading\n\n# Verify numerical correctness\nare_close = jnp.allclose(\n    result_activation[0],    # Result from activation offloading only\n    result_both[0],         # Result from both activation and parameter offloading\n    rtol=1e-5,\n    atol=1e-5\n)\nprint(f\"Results match within tolerance: {are_close}\")\n```\n\nExample:\n```text\nTemp size: 4.75 MB\nArgument size: 0.25 MB\nTotal size: 25.00 MB\nResults match within tolerance: True\n```\n\nExample:\n```text\nimport optax\n\nDIM = 7168\n\n# Initialize data and parameter w1, w2, w3 and w4\ninput = jnp.ones((DIM, DIM))\nparams = {f'w{i}': jnp.ones((DIM, DIM)) for i in range(1, 5)}\n\n# Initialize optimizer\noptimizer = optax.chain(\n    optax.clip_by_global_norm(1.0),\n    optax.adam(learning_rate=0.1)\n)\nopt_state = optimizer.init(params)\n\ndef gelu(x):\n  return 0.5 * x * (1 + jnp.tanh(jnp.sqrt(2 / jnp.pi) * (x + 0.044715 * x**3)))\n\ndef single_layer(x, w):\n  return x @ w\n\ndef forward(params, x):\n  for i in range(1, 5):\n    x = gelu(single_layer(x, params[f'w{i}']))\n  return x\n\ndef compute_loss(params, inputs):\n  outputs = forward(params, inputs)\n  loss = jnp.mean((outputs - inputs) ** 2)\n  l2_reg = 0.001 * sum(jnp.sum(w ** 2) for w in jax.tree_util.tree_leaves(params))\n  return loss + l2_reg\n\ndef step(params, opt_state, inputs):\n  grads = jax.grad(lambda p: compute_loss(p, inputs))(params)\n  updates, new_opt_state = optimizer.update(grads, opt_state, params)\n  return optax.apply_updates(params, updates), new_opt_state\n\n# JIT compile the step function with proper sharding\nstep = jax.jit(step, donate_argnums=(0, 1))\n\n# Run a optimization step\nnew_params, new_opt_state = step(params, opt_state, input)\n\n# Analyze memory usage\ncompiled_step = step.lower(params, opt_state, input).compile()\ncompiled_stats = compiled_step.memory_analysis()\n\nif compiled_stats is not None:\n  total = compiled_stats.temp_size_in_bytes + compiled_stats.argument_size_in_bytes \\\n      + compiled_stats.output_size_in_bytes - compiled_stats.alias_size_in_bytes\n  print(f\"Temp size: {compiled_stats.temp_size_in_bytes / (1024**3):.2f} GB\")\n  print(f\"Argument size: {compiled_stats.argument_size_in_bytes / (1024**3):.2f} GB\")\n  print(f\"Total size: {total / (1024**3):.2f} GB\")\n```\n\nExample:\n```text\nTemp size: 2.11 GB\nArgument size: 2.49 GB\nTotal size: 4.59 GB\n```\n\nExample:\n```text\n# Create sharding specifications for device and host memory\ns_dev = jax.sharding.SingleDeviceSharding(jax.devices()[0], memory_kind=\"device\")\ns_host = jax.sharding.SingleDeviceSharding(jax.devices()[0], memory_kind=\"pinned_host\")\n\ndef step(params, opt_state, inputs):\n  grads = jax.grad(lambda p: compute_loss(p, inputs))(params)\n  opt_state = jax.device_put(opt_state, s_dev)\n  updates, new_opt_state = optimizer.update(grads, opt_state, params)\n  new_params = optax.apply_updates(params, updates)\n  return new_params, new_opt_state\n\nparams = {f'w{i}': jnp.ones((DIM, DIM)) for i in range(1, 5)}\nopt_state = optimizer.init(params)\n\n# Initialize optimizer\noptimizer = optax.chain(\n  optax.clip_by_global_norm(1.0),\n  optax.adam(learning_rate=0.1)\n)\n\n# Optimizer state is placed on the host during initialization\nopt_state = jax.device_put(opt_state, s_host)\n\n# JIT compile the step function with proper sharding and memory optimization\nstep = jax.jit(\n  step,\n  donate_argnums=(0,),\n  out_shardings=(s_dev, s_host)\n)\n\n# Run an optimization step\nnew_params, offload_opt_state = step(params, opt_state, input)\n\n# Analyze memory usage\ncompiled_step = step.lower(params, opt_state, input).compile()\ncompiled_stats = compiled_step.memory_analysis()\nif compiled_stats is not None:\n  total = compiled_stats.temp_size_in_bytes + compiled_stats.argument_size_in_bytes \\\n      + compiled_stats.output_size_in_bytes - compiled_stats.alias_size_in_bytes\n  print(f\"Temp size: {compiled_stats.temp_size_in_bytes / (1024**3):.2f} GB\")\n  print(f\"Argument size: {compiled_stats.argument_size_in_bytes / (1024**3):.2f} MB\")\n  print(f\"Total size: {total / (1024**3):.2f} GB\")\n```\n\nExample:\n```text\nTemp size: 1.91 GB\nArgument size: 0.96 MB\nTotal size: 2.87 GB\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.766Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":22,"totalLines":390,"estimatedTokens":3096}}72{"id":"doc-control_autodiff_s_saved_values_with_jax_checkpo-ee89b42c","source":"documentation","title":"Control autodiff’s saved values with jax.checkpoint (aka jax.remat) — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/autodiff_remat.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n```\n\nExample:\n```text\ndef g(W, x):\n  y = jnp.dot(W, x)\n  return jnp.sin(y)\n\ndef f(W1, W2, W3, x):\n  x = g(W1, x)\n  x = g(W2, x)\n  x = g(W3, x)\n  return x\n\nW1 = jnp.ones((5, 4))\nW2 = jnp.ones((6, 5))\nW3 = jnp.ones((7, 6))\nx = jnp.ones(4)\n\n# Inspect the 'residual' values to be saved on the forward pass\n# if we were to evaluate `jax.grad(f)(W1, W2, W3, x)`\nfrom jax.ad_checkpoint import print_saved_residuals\njax.ad_checkpoint.print_saved_residuals(f, W1, W2, W3, x)\n```\n\nExample:\n```text\nf32[5,4] from the argument 'W1'\nf32[6,5] from the argument 'W2'\nf32[7,6] from the argument 'W3'\nf32[4] from the argument 'x'\nf32[5] output of sin from <ipython-input-4-f510dde58e22>:3 (g)\nf32[5] output of cos from <ipython-input-4-f510dde58e22>:3 (g)\nf32[6] output of sin from <ipython-input-4-f510dde58e22>:3 (g)\nf32[6] output of cos from <ipython-input-4-f510dde58e22>:3 (g)\nf32[7] output of cos from <ipython-input-4-f510dde58e22>:3 (g)\n```\n\nExample:\n```text\ndef f2(W1, W2, W3, x):\n  x = jax.checkpoint(g)(W1, x)\n  x = jax.checkpoint(g)(W2, x)\n  x = jax.checkpoint(g)(W3, x)\n  return x\n\njax.ad_checkpoint.print_saved_residuals(f2, W1, W2, W3, x)\n```\n\nExample:\n```text\nf32[5,4] from the argument 'W1'\nf32[6,5] from the argument 'W2'\nf32[7,6] from the argument 'W3'\nf32[4] from the argument 'x'\nf32[5] output of sin from <ipython-input-4-f510dde58e22>:3 (g)\nf32[6] output of sin from <ipython-input-4-f510dde58e22>:3 (g)\n```\n\nExample:\n```text\nf3 = jax.checkpoint(f, policy=jax.checkpoint_policies.dots_with_no_batch_dims_saveable)\njax.ad_checkpoint.print_saved_residuals(f3, W1, W2, W3, x)\n```\n\nExample:\n```text\nf32[5,4] from the argument 'W1'\nf32[6,5] from the argument 'W2'\nf32[7,6] from the argument 'W3'\nf32[4] from the argument 'x'\nf32[5] output of dot_general from <ipython-input-4-f510dde58e22>:2 (g)\nf32[6] output of dot_general from <ipython-input-4-f510dde58e22>:2 (g)\nf32[7] output of dot_general from <ipython-input-4-f510dde58e22>:2 (g)\n```\n\nExample:\n```text\nfrom jax.ad_checkpoint import checkpoint_name\n\ndef f4(W1, W2, W3, x):\n  x = checkpoint_name(g(W1, x), name='a')\n  x = checkpoint_name(g(W2, x), name='b')\n  x = checkpoint_name(g(W3, x), name='c')\n  return x\n\nf4 = jax.checkpoint(f4, policy=jax.checkpoint_policies.save_only_these_names('a'))\njax.ad_checkpoint.print_saved_residuals(f4, W1, W2, W3, x)\n```\n\nExample:\n```text\nf32[5,4] from the argument 'W1'\nf32[6,5] from the argument 'W2'\nf32[7,6] from the argument 'W3'\nf32[4] from the argument 'x'\nf32[5] named 'a' from <ipython-input-7-fc0ed1c14b8d>:4 (f4)\n```\n\nExample:\n```text\nfrom jax.tree_util import tree_flatten, tree_unflatten\n\nfrom rich.console import Console\nfrom rich.table import Table\nimport rich.text\n\ndef print_fwd_bwd(f, *args, **kwargs) -> None:\n  args, in_tree = tree_flatten((args, kwargs))\n\n  def f_(*args):\n    args, kwargs = tree_unflatten(in_tree, args)\n    return f(*args, **kwargs)\n\n  fwd = jax.make_jaxpr(lambda *args: jax.vjp(f_, *args))(*args).jaxpr\n\n  y, f_vjp = jax.vjp(f_, *args)\n  res, in_tree = tree_flatten(f_vjp)\n\n  def g_(*args):\n    *res, y = args\n    f_vjp = tree_unflatten(in_tree, res)\n    return f_vjp(y)\n\n  bwd = jax.make_jaxpr(g_)(*res, y).jaxpr\n\n  table = Table(show_header=False, show_lines=True, padding=(1, 2, 0, 2), box=None)\n  table.add_row(\"[bold green]forward computation:\",\n                \"[bold green]backward computation:\")\n  table.add_row(rich.text.Text.from_ansi(str(fwd)),\n                rich.text.Text.from_ansi(str(bwd)))\n  console = Console(width=240, force_jupyter=True)\n  console.print(table)\n\ndef _renderable_repr(self):\n  return self.html\nrich.jupyter.JupyterRenderable._repr_html_ = _renderable_repr\n```\n\nExample:\n```text\n# no use of jax.checkpoint:\nprint_fwd_bwd(f, W1, W2, W3, x)\n```\n\nExample:\n```text\nforward computation:                                                        backward computation:                                                                   \n                                                                                                                                                                      \n  { lambda ; a:f32[5,4] b:f32[6,5] c:f32[7,6] d:f32[4]. let                   { lambda ; a:f32[7] b:f32[6] c:f32[7,6] d:f32[6] e:f32[5] f:f32[6,5] g:f32[5] h:f32[4]  \n      e:f32[5] = dot_general[dimension_numbers=(([1], [0]), ([], []))] a d        i:f32[5,4] j:f32[7]. let                                                            \n      f:f32[5] = sin e                                                            k:f32[7] = mul j a                                                                  \n      g:f32[5] = cos e                                                            l:f32[6] = dot_general[dimension_numbers=(([0], [0]), ([], []))] k c                \n      h:f32[6] = dot_general[dimension_numbers=(([1], [0]), ([], []))] b f        m:f32[7,6] = dot_general[dimension_numbers=(([], []), ([], []))] k b                \n      i:f32[6] = sin h                                                            n:f32[6] = mul l d                                                                  \n      j:f32[6] = cos h                                                            o:f32[5] = dot_general[dimension_numbers=(([0], [0]), ([], []))] n f                \n      k:f32[7] = dot_general[dimension_numbers=(([1], [0]), ([], []))] c i        p:f32[6,5] = dot_general[dimension_numbers=(([], []), ([], []))] n e                \n      l:f32[7] = sin k                                                            q:f32[5] = mul o g                                                                  \n      m:f32[7] = cos k                                                            r:f32[4] = dot_general[dimension_numbers=(([0], [0]), ([], []))] q i                \n    in (l, m, i, c, j, f, b, g, d, a) }                                           s:f32[5,4] = dot_general[dimension_numbers=(([], []), ([], []))] q h                \n                                                                                in (s, p, m, r) }\n```\n\nExample:\n```text\n# using jax.checkpoint with policy=jax.checkpoint_policies.dots_with_no_batch_dims_saveable:\nprint_fwd_bwd(f3, W1, W2, W3, x)\n```\n\nExample:\n```text\nforward computation:                                                        backward computation:                                                                          \n                                                                                                                                                                             \n  { lambda ; a:f32[5,4] b:f32[6,5] c:f32[7,6] d:f32[4]. let                   { lambda ; a:f32[5] b:f32[6] c:f32[7] d:f32[5,4] e:f32[6,5] f:f32[7,6] g:f32[4] h:f32[7]. let  \n      e:f32[5] = dot_general[dimension_numbers=(([1], [0]), ([], []))] a d        i:f32[5,4] j:f32[6,5] k:f32[7,6] l:f32[4] = remat2[                                        \n      f:f32[5] = sin e                                                              differentiated=True                                                                      \n      g:f32[6] = dot_general[dimension_numbers=(([1], [0]), ([], []))] b f          jaxpr={ lambda ; m:f32[5] n:f32[6] o:f32[7] p:f32[5,4] q:f32[6,5] r:f32[7,6]             \n      h:f32[6] = sin g                                                                  s:f32[4] t:f32[7]. let                                                               \n      i:f32[7] = dot_general[dimension_numbers=(([1], [0]), ([], []))] c h              u:f32[5] = sin m                                                                     \n      j:f32[7] = sin i                                                                  v:f32[5] = cos m                                                                     \n    in (j, e, g, i, a, b, c, d) }                                                       w:f32[6] = sin n                                                                     \n                                                                                        x:f32[6] = cos n                                                                     \n                                                                                        y:f32[7] = cos o                                                                     \n                                                                                        z:f32[7] = mul t y                                                                   \n                                                                                        ba:f32[6] = dot_general[dimension_numbers=(([0], [0]), ([], []))] z r                \n                                                                                        bb:f32[6] = mul ba x                                                                 \n                                                                                        bc:f32[5] = dot_general[dimension_numbers=(([0], [0]), ([], []))] bb q               \n                                                                                        bd:f32[5] = mul bc v                                                                 \n                                                                                        be:f32[4] = dot_general[dimension_numbers=(([0], [0]), ([], []))] bd p               \n                                                                                        bf:f32[5,4] = dot_general[dimension_numbers=(([], []), ([], []))] bd s               \n                                                                                        bg:f32[6,5] = dot_general[dimension_numbers=(([], []), ([], []))] bb u               \n                                                                                        bh:f32[7,6] = dot_general[dimension_numbers=(([], []), ([], []))] z w                \n                                                                                      in (bf, bg, bh, be) }                                                                  \n                                                                                    policy=<function dot_with_no_batch_dims at 0x7f5e469b1700>                               \n                                                                                    prevent_cse=True                                                                         \n                                                                                  ] a b c d e f g h                                                                          \n                                                                                in (i, j, k, l) }\n```\n\nExample:\n```text\ndef sin_vjp(x):\n  y = jnp.sin(x)\n  cos_x = jnp.cos(x)\n  return y, lambda y_bar: cos_x * y_bar\n```\n\nExample:\n```text\ndef sin_vjp2(x):\n  y = jnp.sin(x)\n  return y, lambda y_bar: jnp.cos(x) * y_bar\n```\n\nExample:\n```text\ndef f(x):\n  y = g(x)\n  z = h(y)\n  return z\n\ndef f_vjp(x):\n  y, g_vjp = jax.vjp(g, x)\n  z, h_vjp = jax.vjp(h, y)\n  def f_bwd(z_bar):\n    y_bar, = h_vjp(z_bar)\n    x_bar, = g_vjp(y_bar)\n    return x_bar\n  return z, f_bwd\n```\n\nExample:\n```text\ndef f_vjp_checkpoint(x):\n  y = g(x)\n  z, h_vjp = jax.vjp(h, y)\n  def f_bwd2(z_bar):\n    y_bar, = h_vjp(z_bar)\n    _, g_vjp = jax.vjp(g, x)\n    x_bar, = g_vjp(y_bar)\n    return x_bar\n  return z, f_bwd2\n```\n\nExample:\n```text\ndef f_checkpoint(x):\n  y = jax.checkpoint(g)(x)\n  z = h(y)\n  return z\n```\n\nExample:\n```text\ndef f_checkpoint_grad(x):\n  y = g(x)                  # step 1\n  _, h_vjp = jax.vjp(h)(y)  # step 2\n  y_bar, = h_vjp(1.0)       # step 3\n  _, g_vjp = jax.vjp(g, x)  # step 4\n  x_bar, = g_vjp(y_bar)     # step 5\n  return x_bar\n```\n\nExample:\n```text\ndef f_grad_bad1(x):\n  _ = f(x)                  # step 1\n  _, f_vjp = jax.vjp(f, x)  # step 2\n  x_bar, = f_vjp(1.0)       # step 3\n  return x_bar\n```\n\nExample:\n```text\ndef f_grad_bad2(x):\n  y, g_vjp = jax.vjp(g, x)  # step 1\n  _z = h(y)                 # step 2\n  _, h_vjp = jax.vjp(h, y)  # step 3\n  y_bar, = h_vjp(1.0)       # step 3\n  x_bar, = g_vjp(y_bar)     # step 5\n  return x_bar\n```\n\nExample:\n```text\ndef loss(params, x, y):\n  return jnp.sum((predict(params, x) - y)**2)\n\ndef predict(params, x):\n  *Ws, Wlast = params\n  for W in Ws:\n    x = layer(W, x)\n  x = jnp.dot(Wlast, x)\n  return x\n\ndef layer(W, x):\n  return jnp.sin(jnp.dot(W, x))\n```\n\nExample:\n```text\nW1 = W2 = W3 = jnp.ones((4, 4))\nparams = [W1, W2, W3]\nx = jnp.ones(4)\ny = jnp.ones(4)\n```\n\nExample:\n```text\nprint_saved_residuals(loss, params, x, y)\n```\n\nExample:\n```text\nf32[4,4] from the argument 'params'\nf32[4,4] from the argument 'params'\nf32[4,4] from the argument 'params'\nf32[4] from the argument 'x'\nf32[4] output of sin from <ipython-input-18-3808b5023c3d>:12 (layer)\nf32[4] output of cos from <ipython-input-18-3808b5023c3d>:12 (layer)\nf32[4] output of sin from <ipython-input-18-3808b5023c3d>:12 (layer)\nf32[4] output of cos from <ipython-input-18-3808b5023c3d>:12 (layer)\nf32[4] output of mul from <ipython-input-18-3808b5023c3d>:2 (loss)\n```\n\nExample:\n```text\nloss_checkpoint = jax.checkpoint(loss, policy=jax.checkpoint_policies.dots_with_no_batch_dims_saveable)\nprint_saved_residuals(loss_checkpoint, params, x, y)\n```\n\nExample:\n```text\nf32[4,4] from the argument 'params'\nf32[4,4] from the argument 'params'\nf32[4,4] from the argument 'params'\nf32[4] from the argument 'x'\nf32[4] from the argument 'y'\nf32[4] output of dot_general from <ipython-input-18-3808b5023c3d>:12 (layer)\nf32[4] output of dot_general from <ipython-input-18-3808b5023c3d>:12 (layer)\nf32[4] output of dot_general from <ipython-input-18-3808b5023c3d>:8 (predict)\n```\n\nExample:\n```text\ndef predict(params, x):\n  *Ws, Wlast = params\n  for i, W in enumerate(Ws):\n    x = layer(W, x)\n    x = checkpoint_name(x, name=f'layer{i}_output')\n  x = jnp.dot(Wlast, x)\n  return x\n```\n\nExample:\n```text\nf32[4,4] from the argument 'params'\nf32[4,4] from the argument 'params'\nf32[4,4] from the argument 'params'\nf32[4] from the argument 'x'\nf32[4] output of cos from <ipython-input-18-3808b5023c3d>:12 (layer)\nf32[4] named 'layer0_output' from <ipython-input-22-e48aedf368ad>:7 (predict)\nf32[4] output of cos from <ipython-input-18-3808b5023c3d>:12 (layer)\nf32[4] named 'layer1_output' from <ipython-input-22-e48aedf368ad>:7 (predict)\nf32[4] output of mul from <ipython-input-18-3808b5023c3d>:2 (loss)\n```\n\nExample:\n```text\nloss_checkpoint2 = jax.checkpoint(loss, policy=jax.checkpoint_policies.save_any_names_but_these('layer1_output'))\nprint_saved_residuals(loss_checkpoint2, params, x, y)\n```\n\nExample:\n```text\nf32[4,4] from the argument 'params'\nf32[4,4] from the argument 'params'\nf32[4,4] from the argument 'params'\nf32[4] from the argument 'x'\nf32[4] from the argument 'y'\n```\n\nExample:\n```text\ndef chain_compose(funs):\n  def f(x):\n    for fun in funs:\n      x = fun(x)\n    return x\n  return f\n\nf = chain_compose([jnp.sin] * 8)\nprint_saved_residuals(f, 3.)\n```\n\nExample:\n```text\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\n```\n\nExample:\n```text\nf = chain_compose([jnp.sin] * 16)\nprint_saved_residuals(f, 3.)\n```\n\nExample:\n```text\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\nf32[] output of cos from <ipython-input-25-46b5594773cb>:4 (f)\n```\n\nExample:\n```text\ndef recursive_checkpoint(funs):\n  if len(funs) == 1:\n    return funs[0]\n  elif len(funs) == 2:\n    f1, f2 = funs\n    return lambda x: f1(f2(x))\n  else:\n    f1 = recursive_checkpoint(funs[:len(funs)//2])\n    f2 = recursive_checkpoint(funs[len(funs)//2:])\n    return lambda x: f1(jax.checkpoint(f2)(x))\n```\n\nExample:\n```text\nf = recursive_checkpoint([jnp.sin] * 8)\nprint_saved_residuals(f, 3.)\n```\n\nExample:\n```text\nf32[] from the argument 'x'\nf32[] output of sin from <ipython-input-27-86f83c871e81>:6 (<lambda>)\nf32[] output of cos from <ipython-input-27-86f83c871e81>:6 (<lambda>)\nf32[] output of cos from <ipython-input-27-86f83c871e81>:6 (<lambda>)\n```\n\nExample:\n```text\nf = recursive_checkpoint([jnp.sin] * 16)\nprint_saved_residuals(f, 3.)\n```\n\nExample:\n```text\nf32[] from the argument 'x'\nf32[] output of sin from <ipython-input-27-86f83c871e81>:6 (<lambda>)\nf32[] output of sin from <ipython-input-27-86f83c871e81>:6 (<lambda>)\nf32[] output of cos from <ipython-input-27-86f83c871e81>:6 (<lambda>)\nf32[] output of cos from <ipython-input-27-86f83c871e81>:6 (<lambda>)\n```\n\nExample:\n```text\nf = chain_compose([jnp.sin] * 8)\nprint_fwd_bwd(f, 3.)\n```\n\nExample:\n```text\nforward computation:                  backward computation:                                                                    \n                                                                                                                                 \n  { lambda ; a:f32[]. let               { lambda ; a:f32[] b:f32[] c:f32[] d:f32[] e:f32[] f:f32[] g:f32[] h:f32[] i:f32[]. let  \n      b:f32[] = sin a                       j:f32[] = mul i a                                                                    \n      c:f32[] = cos a                       k:f32[] = mul j b                                                                    \n      d:f32[] = sin b                       l:f32[] = mul k c                                                                    \n      e:f32[] = cos b                       m:f32[] = mul l d                                                                    \n      f:f32[] = sin d                       n:f32[] = mul m e                                                                    \n      g:f32[] = cos d                       o:f32[] = mul n f                                                                    \n      h:f32[] = sin f                       p:f32[] = mul o g                                                                    \n      i:f32[] = cos f                       q:f32[] = mul p h                                                                    \n      j:f32[] = sin h                     in (q,) }                                                                              \n      k:f32[] = cos h                                                                                                            \n      l:f32[] = sin j                                                                                                            \n      m:f32[] = cos j                                                                                                            \n      n:f32[] = sin l                                                                                                            \n      o:f32[] = cos l                                                                                                            \n      p:f32[] = sin n                                                                                                            \n      q:f32[] = cos n                                                                                                            \n    in (p, q, o, m, k, i, g, e, c) }\n```\n\nExample:\n```text\nf = recursive_checkpoint([jnp.sin] * 8)\nprint_fwd_bwd(f, 3.)\n```\n\nExample:\n```text\nforward computation:                                                              backward computation:                               \n                                                                                                                                        \n  { lambda ; a:f32[]. let                                                           { lambda ; a:f32[] b:f32[] c:f32[] d:f32[]. let     \n      b:f32[] = remat2[                                                                 e:f32[] = mul d a                               \n        differentiated=False                                                            f:f32[] = mul e b                               \n        jaxpr={ lambda ; c:f32[]. let d:f32[] = sin c; e:f32[] = sin d in (e,) }        g:f32[] = remat2[                               \n        policy=None                                                                       differentiated=True                           \n        prevent_cse=True                                                                  jaxpr={ lambda ; h:f32[] i:f32[]. let         \n      ] a                                                                                     j:f32[] = sin h                           \n      f:f32[] = sin b                                                                         k:f32[] = cos h                           \n      g:f32[] = sin f                                                                         l:f32[] = cos j                           \n      h:f32[] = sin g                                                                         m:f32[] = mul i l                         \n      i:f32[] = sin h                                                                         n:f32[] = mul m k                         \n      j:f32[] = sin i                                                                       in (n,) }                                   \n      k:f32[] = cos i                                                                     policy=None                                   \n      l:f32[] = sin j                                                                     prevent_cse=True                              \n      m:f32[] = cos j                                                                   ] c f                                           \n    in (l, m, k, g, a) }                                                                o:f32[] = remat2[                               \n                                                                                          differentiated=True                           \n                                                                                          jaxpr={ lambda ; p:f32[] q:f32[]. let         \n                                                                                              r:f32[] = sin p                           \n                                                                                              s:f32[] = sin r                           \n                                                                                              t:f32[] = sin s                           \n                                                                                              u:f32[] = cos s                           \n                                                                                              v:f32[] = cos t                           \n                                                                                              w:f32[] = mul q v                         \n                                                                                              x:f32[] = mul w u                         \n                                                                                              y:f32[] = remat2[                         \n                                                                                                differentiated=True                     \n                                                                                                jaxpr={ lambda ; z:f32[] ba:f32[]. let  \n                                                                                                    bb:f32[] = sin z                    \n                                                                                                    bc:f32[] = cos z                    \n                                                                                                    bd:f32[] = cos bb                   \n                                                                                                    be:f32[] = mul ba bd                \n                                                                                                    bf:f32[] = mul be bc                \n                                                                                                  in (bf,) }                            \n                                                                                                policy=None                             \n                                                                                                prevent_cse=True                        \n                                                                                              ] p x                                     \n                                                                                            in (y,) }                                   \n                                                                                          policy=None                                   \n                                                                                          prevent_cse=True                              \n                                                                                        ] 3.0 g                                         \n                                                                                      in (o,) }\n```\n\nExample:\n```text\nLayerParam = tuple[jnp.ndarray, jnp.ndarray]  # weights, bias pair for a layer\nParamsList = list[LayerParam]\n\ndef net(params: ParamsList, x: jnp.ndarray):\n  for W, b in params:\n    x = jnp.maximum(jnp.dot(x, W) + b, 0.)\n  return x\n```\n\nExample:\n```text\nStackedWeights = jnp.ndarray  # all weight matrices stacked together\nStackedBiases = jnp.ndarray   # all bias vectors stacked together\n\nall_weights = jnp.stack([W for W, _ in params])\nall_biases = jnp.stack([b for _, b in params])\n\ndef layer(x, W_b_pair):\n  W, b = W_b_pair\n  out = jnp.maximum(jnp.dot(x, W) + b, 0.)\n  return out, None\n\ndef net(all_weights, all_biases, x):\n  x, _ = jax.lax.scan(layer, x, (all_weights, all_biases))\n  return x\n```\n\nExample:\n```text\nfrom functools import partial\n\n@partial(jax.checkpoint,\n         policy=jax.checkpoint_policies.dots_with_no_batch_dims_saveable)\ndef layer(x, W_b_pair):\n  W, b = W_b_pair\n  out = jnp.maximum(jnp.dot(x, W) + b, 0.)\n  return out, None\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.768Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":48,"totalLines":606,"estimatedTokens":6835}}73{"id":"doc-custom_derivative_rules_for_jax_transformable_py-79af3aff","source":"documentation","title":"Custom derivative rules for JAX-transformable Python functions — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html","text":"Example:\n```text\nimport jax.numpy as jnp\nfrom jax import custom_jvp\n\n@custom_jvp\ndef f(x, y):\n  return jnp.sin(x) * y\n\n@f.defjvp\ndef f_jvp(primals, tangents):\n  x, y = primals\n  x_dot, y_dot = tangents\n  primal_out = f(x, y)\n  tangent_out = jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot\n  return primal_out, tangent_out\n```\n\nExample:\n```text\nfrom jax import jvp, grad\n\nprint(f(2., 3.))\ny, y_dot = jvp(f, (2., 3.), (1., 0.))\nprint(y)\nprint(y_dot)\nprint(grad(f)(2., 3.))\n```\n\nExample:\n```text\n2.7278922\n2.7278922\n-1.2484405\n-1.2484405\n```\n\nExample:\n```text\n# Equivalent alternative using the defjvps convenience wrapper\n\n@custom_jvp\ndef f(x, y):\n  return jnp.sin(x) * y\n\nf.defjvps(lambda x_dot, primal_out, x, y: jnp.cos(x) * x_dot * y,\n          lambda y_dot, primal_out, x, y: jnp.sin(x) * y_dot)\n```\n\nExample:\n```text\nprint(f(2., 3.))\ny, y_dot = jvp(f, (2., 3.), (1., 0.))\nprint(y)\nprint(y_dot)\nprint(grad(f)(2., 3.))\n```\n\nExample:\n```text\nfrom jax import custom_vjp\n\n@custom_vjp\ndef f(x, y):\n  return jnp.sin(x) * y\n\ndef f_fwd(x, y):\n  # Returns primal output and residuals to be used in backward pass by f_bwd.\n  return f(x, y), (jnp.cos(x), jnp.sin(x), y)\n\ndef f_bwd(res, g):\n  cos_x, sin_x, y = res # Gets residuals computed in f_fwd\n  return (cos_x * g * y, sin_x * g)\n\nf.defvjp(f_fwd, f_bwd)\n```\n\nExample:\n```text\nprint(grad(f)(2., 3.))\n```\n\nExample:\n```text\n-1.2484405\n```\n\nExample:\n```text\ndef log1pexp(x):\n  return jnp.log(1. + jnp.exp(x))\n\nlog1pexp(3.)\n```\n\nExample:\n```text\nArray(3.0485873, dtype=float32, weak_type=True)\n```\n\nExample:\n```text\nfrom jax import jit, grad, vmap\n\nprint(jit(log1pexp)(3.))\nprint(jit(grad(log1pexp))(3.))\nprint(vmap(jit(grad(log1pexp)))(jnp.arange(3.)))\n```\n\nExample:\n```text\n3.0485873\n0.95257413\n[0.5       0.7310586 0.8807971]\n```\n\nExample:\n```text\nprint(grad(log1pexp)(100.))\n```\n\nExample:\n```text\nnan\n```\n\nExample:\n```text\nfrom jax import make_jaxpr\n\nmake_jaxpr(grad(log1pexp))(100.)\n```\n\nExample:\n```text\n{ lambda ; a:f32[]. let\n    b:f32[] = exp a\n    c:f32[] = add 1.0:f32[] b\n    _:f32[] = log c\n    d:f32[] = div 1.0:f32[] c\n    e:f32[] = mul d b\n  in (e,) }\n```\n\nExample:\n```text\nfrom jax import custom_jvp\n\n@custom_jvp\ndef log1pexp(x):\n  return jnp.log(1. + jnp.exp(x))\n\n@log1pexp.defjvp\ndef log1pexp_jvp(primals, tangents):\n  x, = primals\n  x_dot, = tangents\n  ans = log1pexp(x)\n  ans_dot = (1 - 1/(1 + jnp.exp(x))) * x_dot\n  return ans, ans_dot\n```\n\nExample:\n```text\n1.0\n```\n\nExample:\n```text\nprint(jit(log1pexp)(3.))\nprint(jit(grad(log1pexp))(3.))\nprint(vmap(jit(grad(log1pexp)))(jnp.arange(3.)))\n```\n\nExample:\n```text\n@custom_jvp\ndef log1pexp(x):\n  return jnp.log(1. + jnp.exp(x))\n\nlog1pexp.defjvps(lambda t, ans, x: (1 - 1/(1 + jnp.exp(x))) * t)\n```\n\nExample:\n```text\nprint(grad(log1pexp)(100.))\nprint(jit(log1pexp)(3.))\nprint(jit(grad(log1pexp))(3.))\nprint(vmap(jit(grad(log1pexp)))(jnp.arange(3.)))\n```\n\nExample:\n```text\n1.0\n3.0485873\n0.95257413\n[0.5       0.7310586 0.8807971]\n```\n\nExample:\n```text\ndef f(x):\n  return x / (1 + jnp.sqrt(x))\n```\n\nExample:\n```text\nprint(grad(f)(0.))\n```\n\nExample:\n```text\n@custom_jvp\ndef f(x):\n  return x / (1 + jnp.sqrt(x))\n\n@f.defjvp\ndef f_jvp(primals, tangents):\n  x, = primals\n  x_dot, = tangents\n  ans = f(x)\n  ans_dot = ((jnp.sqrt(x) + 2) / (2 * (jnp.sqrt(x) + 1)**2)) * x_dot\n  return ans, ans_dot\n```\n\nExample:\n```text\n@custom_jvp\ndef f(x):\n  return x / (1 + jnp.sqrt(x))\n\nf.defjvps(lambda t, ans, x: ((jnp.sqrt(x) + 2) / (2 * (jnp.sqrt(x) + 1)**2)) * t)\n```\n\nExample:\n```text\nfrom functools import partial\nfrom jax import custom_vjp\n\n@custom_vjp\ndef clip_gradient(lo, hi, x):\n  return x  # identity function\n\ndef clip_gradient_fwd(lo, hi, x):\n  return x, (lo, hi)  # save bounds as residuals\n\ndef clip_gradient_bwd(res, g):\n  lo, hi = res\n  return (None, None, jnp.clip(g, lo, hi))  # use None to indicate zero cotangents for lo and hi\n\nclip_gradient.defvjp(clip_gradient_fwd, clip_gradient_bwd)\n```\n\nExample:\n```text\nimport matplotlib.pyplot as plt\nfrom jax import vmap\n\nt = jnp.linspace(0, 10, 1000)\n\nplt.plot(jnp.sin(t))\nplt.plot(vmap(grad(jnp.sin))(t))\n```\n\nExample:\n```text\n[<matplotlib.lines.Line2D at 0x7fe87fcc96a0>]\n```\n\nExample:\n```text\ndef clip_sin(x):\n  x = clip_gradient(-0.75, 0.75, x)\n  return jnp.sin(x)\n\nplt.plot(clip_sin(t))\nplt.plot(vmap(grad(clip_sin))(t))\n```\n\nExample:\n```text\n[<matplotlib.lines.Line2D at 0x7fe87faa7f80>]\n```\n\nExample:\n```text\nfrom jax.lax import while_loop\n\ndef fixed_point(f, a, x_guess):\n  def cond_fun(carry):\n    x_prev, x = carry\n    return jnp.abs(x_prev - x) > 1e-6\n\n  def body_fun(carry):\n    _, x = carry\n    return x, f(a, x)\n\n  _, x_star = while_loop(cond_fun, body_fun, (x_guess, f(a, x_guess)))\n  return x_star\n```\n\nExample:\n```text\ndef newton_sqrt(a):\n  update = lambda a, x: 0.5 * (x + a / x)\n  return fixed_point(update, a, a)\n```\n\nExample:\n```text\nprint(newton_sqrt(2.))\n```\n\nExample:\n```text\n1.4142135\n```\n\nExample:\n```text\nprint(jit(vmap(newton_sqrt))(jnp.array([1., 2., 3., 4.])))\n```\n\nExample:\n```text\n[1.        1.4142135 1.7320509 2.       ]\n```\n\nExample:\n```text\nfrom jax import vjp\n\n@partial(custom_vjp, nondiff_argnums=(0,))\ndef fixed_point(f, a, x_guess):\n  def cond_fun(carry):\n    x_prev, x = carry\n    return jnp.abs(x_prev - x) > 1e-6\n\n  def body_fun(carry):\n    _, x = carry\n    return x, f(a, x)\n\n  _, x_star = while_loop(cond_fun, body_fun, (x_guess, f(a, x_guess)))\n  return x_star\n\ndef fixed_point_fwd(f, a, x_init):\n  x_star = fixed_point(f, a, x_init)\n  return x_star, (a, x_star)\n\ndef fixed_point_rev(f, res, x_star_bar):\n  a, x_star = res\n  _, vjp_a = vjp(lambda a: f(a, x_star), a)\n  a_bar, = vjp_a(fixed_point(partial(rev_iter, f),\n                             (a, x_star, x_star_bar),\n                             x_star_bar))\n  return a_bar, jnp.zeros_like(x_star)\n\ndef rev_iter(f, packed, u):\n  a, x_star, x_star_bar = packed\n  _, vjp_x = vjp(lambda x: f(a, x), x_star)\n  return x_star_bar + vjp_x(u)[0]\n\nfixed_point.defvjp(fixed_point_fwd, fixed_point_rev)\n```\n\nExample:\n```text\nprint(grad(newton_sqrt)(2.))\nprint(grad(grad(newton_sqrt))(2.))\n```\n\nExample:\n```text\n0.35355338\n-0.088388346\n```\n\nExample:\n```text\nprint(grad(jnp.sqrt)(2.))\nprint(grad(grad(jnp.sqrt))(2.))\n```\n\nExample:\n```text\n0.35355338\n-0.08838835\n```\n\nExample:\n```text\nfrom jax import custom_jvp\nimport jax.numpy as jnp\n\n# f :: a -> b\n@custom_jvp\ndef f(x):\n  return jnp.sin(x)\n\n# f_jvp :: (a, T a) -> (b, T b)\ndef f_jvp(primals, tangents):\n  x, = primals\n  t, = tangents\n  return f(x), jnp.cos(x) * t\n\nf.defjvp(f_jvp)\n```\n\nExample:\n```text\n<function __main__.f_jvp(primals, tangents)>\n```\n\nExample:\n```text\nfrom jax import jvp\n\nprint(f(3.))\n\ny, y_dot = jvp(f, (3.,), (1.,))\nprint(y)\nprint(y_dot)\n```\n\nExample:\n```text\n0.14112\n0.14112\n-0.9899925\n```\n\nExample:\n```text\n@custom_jvp\ndef f(x):\n  ...\n\n@f.defjvp\ndef f_jvp(primals, tangents):\n  ...\n```\n\nExample:\n```text\nfrom jax import grad\n\nprint(grad(f)(3.))\nprint(grad(grad(f))(3.))\n```\n\nExample:\n```text\n-0.9899925\n-0.14112\n```\n\nExample:\n```text\n@custom_jvp\ndef f(x, y):\n  return x ** 2 * y\n\n@f.defjvp\ndef f_jvp(primals, tangents):\n  x, y = primals\n  x_dot, y_dot = tangents\n  primal_out = f(x, y)\n  tangent_out = 2 * x * y * x_dot + x ** 2 * y_dot\n  return primal_out, tangent_out\n```\n\nExample:\n```text\n12.0\n```\n\nExample:\n```text\n@custom_jvp\ndef f(x):\n  return jnp.sin(x)\n\nf.defjvps(lambda t, ans, x: jnp.cos(x) * t)\n```\n\nExample:\n```text\nprint(grad(f)(3.))\n```\n\nExample:\n```text\n-0.9899925\n```\n\nExample:\n```text\n@custom_jvp\ndef f(x, y):\n  return x ** 2 * y\n\nf.defjvps(lambda x_dot, primal_out, x, y: 2 * x * y * x_dot,\n          lambda y_dot, primal_out, x, y: x ** 2 * y_dot)\n```\n\nExample:\n```text\nprint(grad(f)(2., 3.))\nprint(grad(f, 0)(2., 3.))  # same as above\nprint(grad(f, 1)(2., 3.))\n```\n\nExample:\n```text\n12.0\n12.0\n4.0\n```\n\nExample:\n```text\n@custom_jvp\ndef f(x, y):\n  return x ** 2 * y\n\nf.defjvps(lambda x_dot, primal_out, x, y: 2 * x * y * x_dot,\n          None)\n```\n\nExample:\n```text\n12.0\n12.0\n0.0\n```\n\nExample:\n```text\n@custom_jvp\ndef f(x):\n  print('called f!')  # a harmless side-effect\n  return jnp.sin(x)\n\n@f.defjvp\ndef f_jvp(primals, tangents):\n  print('called f_jvp!')  # a harmless side-effect\n  x, = primals\n  t, = tangents\n  return f(x), jnp.cos(x) * t\n```\n\nExample:\n```text\nfrom jax import vmap, jit\n\nprint(f(3.))\n```\n\nExample:\n```text\ncalled f!\n0.14112\n```\n\nExample:\n```text\nprint(vmap(f)(jnp.arange(3.)))\nprint(jit(f)(3.))\n```\n\nExample:\n```text\ncalled f!\n[0.         0.84147096 0.9092974 ]\ncalled f!\n0.14112\n```\n\nExample:\n```text\ny, y_dot = jvp(f, (3.,), (1.,))\nprint(y_dot)\n```\n\nExample:\n```text\ncalled f_jvp!\ncalled f!\n-0.9899925\n```\n\nExample:\n```text\ngrad(grad(f))(3.)\n```\n\nExample:\n```text\ncalled f_jvp!\ncalled f_jvp!\ncalled f!\n```\n\nExample:\n```text\nArray(-0.14112, dtype=float32, weak_type=True)\n```\n\nExample:\n```text\n@custom_jvp\ndef f(x):\n  if x > 0:\n    return jnp.sin(x)\n  else:\n    return jnp.cos(x)\n\n@f.defjvp\ndef f_jvp(primals, tangents):\n  x, = primals\n  x_dot, = tangents\n  ans = f(x)\n  if x > 0:\n    return ans, 2 * x_dot\n  else:\n    return ans, 3 * x_dot\n```\n\nExample:\n```text\nprint(grad(f)(1.))\nprint(grad(f)(-1.))\n```\n\nExample:\n```text\n2.0\n3.0\n```\n\nExample:\n```text\nfrom jax import custom_vjp\nimport jax.numpy as jnp\n\n# f :: a -> b\n@custom_vjp\ndef f(x):\n  return jnp.sin(x)\n\n# f_fwd :: a -> (b, c)\ndef f_fwd(x):\n  return f(x), jnp.cos(x)\n\n# f_bwd :: (c, CT b) -> CT a\ndef f_bwd(cos_x, y_bar):\n  return (cos_x * y_bar,)\n\nf.defvjp(f_fwd, f_bwd)\n```\n\nExample:\n```text\nfrom jax import grad\n\nprint(f(3.))\nprint(grad(f)(3.))\n```\n\nExample:\n```text\n0.14112\n-0.9899925\n```\n\nExample:\n```text\nfrom jax import custom_vjp\n\n@custom_vjp\ndef f(x, y):\n  return jnp.sin(x) * y\n\ndef f_fwd(x, y):\n  return f(x, y), (jnp.cos(x), jnp.sin(x), y)\n\ndef f_bwd(res, g):\n  cos_x, sin_x, y = res\n  return (cos_x * g * y, sin_x * g)\n\nf.defvjp(f_fwd, f_bwd)\n```\n\nExample:\n```text\n@custom_vjp\ndef f(x):\n  print(\"called f!\")\n  return jnp.sin(x)\n\ndef f_fwd(x):\n  print(\"called f_fwd!\")\n  return f(x), jnp.cos(x)\n\ndef f_bwd(cos_x, y_bar):\n  print(\"called f_bwd!\")\n  return (cos_x * y_bar,)\n\nf.defvjp(f_fwd, f_bwd)\n```\n\nExample:\n```text\nprint(f(3.))\n```\n\nExample:\n```text\ncalled f_fwd!\ncalled f!\ncalled f_bwd!\n-0.9899925\n```\n\nExample:\n```text\ny, f_vjp = vjp(f, 3.)\nprint(y)\n```\n\nExample:\n```text\ncalled f_fwd!\ncalled f!\n0.14112\n```\n\nExample:\n```text\nprint(f_vjp(1.))\n```\n\nExample:\n```text\ncalled f_bwd!\n(Array(-0.9899925, dtype=float32, weak_type=True),)\n```\n\nExample:\n```text\nfrom jax import jvp\n\ntry:\n  jvp(f, (3.,), (1.,))\nexcept TypeError as e:\n  print('ERROR! {}'.format(e))\n```\n\nExample:\n```text\ncalled f_fwd!\ncalled f!\nERROR! can't apply forward-mode autodiff (jvp) to a custom_vjp function.\n```\n\nExample:\n```text\nimport pdb\n\n@custom_vjp\ndef debug(x):\n  return x  # acts like identity\n\ndef debug_fwd(x):\n  return x, x\n\ndef debug_bwd(x, g):\n  pdb.set_trace()\n  return g\n\ndebug.defvjp(debug_fwd, debug_bwd)\n```\n\nExample:\n```text\ndef foo(x):\n  y = x ** 2\n  y = debug(y)  # insert pdb in corresponding backward pass step\n  return jnp.sin(y)\n```\n\nExample:\n```text\njax.grad(foo)(3.)\n\n> <ipython-input-113-b19a2dc1abf7>(12)debug_bwd()\n-> return g\n(Pdb) p x\nArray(9., dtype=float32)\n(Pdb) p g\nArray(-0.91113025, dtype=float32)\n(Pdb) q\n```\n\nExample:\n```text\nfrom collections import namedtuple\nPoint = namedtuple(\"Point\", [\"x\", \"y\"])\n\n@custom_jvp\ndef f(pt):\n  x, y = pt.x, pt.y\n  return {'a': x ** 2,\n          'b': (jnp.sin(x), jnp.cos(y))}\n\n@f.defjvp\ndef f_jvp(primals, tangents):\n  pt, = primals\n  pt_dot, =  tangents\n  ans = f(pt)\n  ans_dot = {'a': 2 * pt.x * pt_dot.x,\n             'b': (jnp.cos(pt.x) * pt_dot.x, -jnp.sin(pt.y) * pt_dot.y)}\n  return ans, ans_dot\n\ndef fun(pt):\n  dct = f(pt)\n  return dct['a'] + dct['b'][0]\n```\n\nExample:\n```text\npt = Point(1., 2.)\n\nprint(f(pt))\n```\n\nExample:\n```text\n{'a': 1.0, 'b': (Array(0.84147096, dtype=float32, weak_type=True), Array(-0.41614684, dtype=float32, weak_type=True))}\n```\n\nExample:\n```text\nprint(grad(fun)(pt))\n```\n\nExample:\n```text\nPoint(x=Array(2.5403023, dtype=float32, weak_type=True), y=Array(0., dtype=float32, weak_type=True))\n```\n\nExample:\n```text\n@custom_vjp\ndef f(pt):\n  x, y = pt.x, pt.y\n  return {'a': x ** 2,\n          'b': (jnp.sin(x), jnp.cos(y))}\n\ndef f_fwd(pt):\n  return f(pt), pt\n\ndef f_bwd(pt, g):\n  a_bar, (b0_bar, b1_bar) = g['a'], g['b']\n  x_bar = 2 * pt.x * a_bar + jnp.cos(pt.x) * b0_bar\n  y_bar = -jnp.sin(pt.y) * b1_bar\n  return (Point(x_bar, y_bar),)\n\nf.defvjp(f_fwd, f_bwd)\n\ndef fun(pt):\n  dct = f(pt)\n  return dct['a'] + dct['b'][0]\n```\n\nExample:\n```text\nPoint(x=Array(2.5403023, dtype=float32, weak_type=True), y=Array(-0., dtype=float32, weak_type=True))\n```\n\nExample:\n```text\nfrom functools import partial\n\n@partial(custom_jvp, nondiff_argnums=(0,))\ndef app(f, x):\n  return f(x)\n\n@app.defjvp\ndef app_jvp(f, primals, tangents):\n  x, = primals\n  x_dot, = tangents\n  return f(x), 2. * x_dot\n```\n\nExample:\n```text\nprint(app(lambda x: x ** 3, 3.))\n```\n\nExample:\n```text\n27.0\n```\n\nExample:\n```text\nprint(grad(app, 1)(lambda x: x ** 3, 3.))\n```\n\nExample:\n```text\n2.0\n```\n\nExample:\n```text\n@partial(custom_jvp, nondiff_argnums=(0, 2))\ndef app2(f, x, g):\n  return f(g((x)))\n\n@app2.defjvp\ndef app2_jvp(f, g, primals, tangents):\n  x, = primals\n  x_dot, = tangents\n  return f(g(x)), 3. * x_dot\n```\n\nExample:\n```text\nprint(app2(lambda x: x ** 3, 3., lambda y: 5 * y))\n```\n\nExample:\n```text\n3375.0\n```\n\nExample:\n```text\nprint(grad(app2, 1)(lambda x: x ** 3, 3., lambda y: 5 * y))\n```\n\nExample:\n```text\n3.0\n```\n\nExample:\n```text\n@partial(custom_vjp, nondiff_argnums=(0,))\ndef app(f, x):\n  return f(x)\n\ndef app_fwd(f, x):\n  return f(x), x\n\ndef app_bwd(f, x, g):\n  return (5 * g,)\n\napp.defvjp(app_fwd, app_bwd)\n```\n\nExample:\n```text\nprint(app(lambda x: x ** 2, 4.))\n```\n\nExample:\n```text\n16.0\n```\n\nExample:\n```text\nprint(grad(app, 1)(lambda x: x ** 2, 4.))\n```\n\nExample:\n```text\n5.0\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.770Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":110,"totalLines":983,"estimatedTokens":3444}}74{"id":"doc-the_autodiff_cookbook_jax_documentation-96e0cf16","source":"documentation","title":"The Autodiff Cookbook — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/autodiff_cookbook.html","text":"Example:\n```text\nimport jax.numpy as jnp\nfrom jax import grad, jit, vmap\nfrom jax import random\n\nkey = random.key(0)\n```\n\nExample:\n```text\ngrad_tanh = grad(jnp.tanh)\nprint(grad_tanh(2.0))\n```\n\nExample:\n```text\n0.070650816\n```\n\nExample:\n```text\nprint(grad(grad(jnp.tanh))(2.0))\nprint(grad(grad(grad(jnp.tanh)))(2.0))\n```\n\nExample:\n```text\n-0.13621868\n0.25265405\n```\n\nExample:\n```text\ndef sigmoid(x):\n    return 0.5 * (jnp.tanh(x / 2) + 1)\n\n# Outputs probability of a label being true.\ndef predict(W, b, inputs):\n    return sigmoid(jnp.dot(inputs, W) + b)\n\n# Build a toy dataset.\ninputs = jnp.array([[0.52, 1.12,  0.77],\n                   [0.88, -1.08, 0.15],\n                   [0.52, 0.06, -1.30],\n                   [0.74, -2.49, 1.39]])\ntargets = jnp.array([True, True, False, True])\n\n# Training loss is the negative log-likelihood of the training examples.\ndef loss(W, b):\n    preds = predict(W, b, inputs)\n    label_probs = preds * targets + (1 - preds) * (1 - targets)\n    return -jnp.sum(jnp.log(label_probs))\n\n# Initialize random model coefficients\nkey, W_key, b_key = random.split(key, 3)\nW = random.normal(W_key, (3,))\nb = random.normal(b_key, ())\n```\n\nExample:\n```text\n# Differentiate `loss` with respect to the first positional argument:\nW_grad = grad(loss, argnums=0)(W, b)\nprint('W_grad', W_grad)\n\n# Since argnums=0 is the default, this does the same thing:\nW_grad = grad(loss)(W, b)\nprint('W_grad', W_grad)\n\n# But we can choose different values too, and drop the keyword:\nb_grad = grad(loss, 1)(W, b)\nprint('b_grad', b_grad)\n\n# Including tuple values\nW_grad, b_grad = grad(loss, (0, 1))(W, b)\nprint('W_grad', W_grad)\nprint('b_grad', b_grad)\n```\n\nExample:\n```text\nW_grad [-0.433146  -0.7354605 -1.2598922]\nW_grad [-0.433146  -0.7354605 -1.2598922]\nb_grad -0.69001776\nW_grad [-0.433146  -0.7354605 -1.2598922]\nb_grad -0.69001776\n```\n\nExample:\n```text\ndef loss2(params_dict):\n    preds = predict(params_dict['W'], params_dict['b'], inputs)\n    label_probs = preds * targets + (1 - preds) * (1 - targets)\n    return -jnp.sum(jnp.log(label_probs))\n\nprint(grad(loss2)({'W': W, 'b': b}))\n```\n\nExample:\n```text\n{'W': Array([-0.433146 , -0.7354605, -1.2598922], dtype=float32), 'b': Array(-0.69001776, dtype=float32)}\n```\n\nExample:\n```text\nfrom jax import value_and_grad\nloss_value, Wb_grad = value_and_grad(loss, (0, 1))(W, b)\nprint('loss value', loss_value)\nprint('loss value', loss(W, b))\n```\n\nExample:\n```text\nloss value 2.9729187\nloss value 2.9729187\n```\n\nExample:\n```text\n# Set a step size for finite differences calculations\neps = 1e-4\n\n# Check b_grad with scalar finite differences\nb_grad_numerical = (loss(W, b + eps / 2.) - loss(W, b - eps / 2.)) / eps\nprint('b_grad_numerical', b_grad_numerical)\nprint('b_grad_autodiff', grad(loss, 1)(W, b))\n\n# Check W_grad with finite differences in a random direction\nkey, subkey = random.split(key)\nvec = random.normal(subkey, W.shape)\nunitvec = vec / jnp.sqrt(jnp.vdot(vec, vec))\nW_grad_numerical = (loss(W + eps / 2. * unitvec, b) - loss(W - eps / 2. * unitvec, b)) / eps\nprint('W_dirderiv_numerical', W_grad_numerical)\nprint('W_dirderiv_autodiff', jnp.vdot(grad(loss)(W, b), unitvec))\n```\n\nExample:\n```text\nb_grad_numerical -0.6890297\nb_grad_autodiff -0.69001776\nW_dirderiv_numerical 1.3041496\nW_dirderiv_autodiff 1.3006744\n```\n\nExample:\n```text\nfrom jax.test_util import check_grads\ncheck_grads(loss, (W, b), order=2)  # check up to 2nd order derivatives\n```\n\nExample:\n```text\ndef hvp(f, x, v):\n    return grad(lambda x: jnp.vdot(grad(f)(x), v))(x)\n```\n\nExample:\n```text\nfrom jax import jacfwd, jacrev\n\n# Isolate the function from the weight matrix to the predictions\nf = lambda W: predict(W, b, inputs)\n\nJ = jacfwd(f)(W)\nprint(\"jacfwd result, with shape\", J.shape)\nprint(J)\n\nJ = jacrev(f)(W)\nprint(\"jacrev result, with shape\", J.shape)\nprint(J)\n```\n\nExample:\n```text\njacfwd result, with shape (4, 3)\n[[ 0.05069415  0.1091874   0.07506633]\n [ 0.14170025 -0.17390487  0.02415345]\n [ 0.12579198  0.01451446 -0.31447992]\n [ 0.00574409 -0.0193281   0.01078958]]\njacrev result, with shape (4, 3)\n[[ 0.05069415  0.10918741  0.07506634]\n [ 0.14170025 -0.17390487  0.02415345]\n [ 0.12579198  0.01451446 -0.31447995]\n [ 0.00574409 -0.0193281   0.01078958]]\n```\n\nExample:\n```text\ndef predict_dict(params, inputs):\n    return predict(params['W'], params['b'], inputs)\n\nJ_dict = jacrev(predict_dict)({'W': W, 'b': b}, inputs)\nfor k, v in J_dict.items():\n    print(\"Jacobian from {} to logits is\".format(k))\n    print(v)\n```\n\nExample:\n```text\nJacobian from W to logits is\n[[ 0.05069415  0.10918741  0.07506634]\n [ 0.14170025 -0.17390487  0.02415345]\n [ 0.12579198  0.01451446 -0.31447995]\n [ 0.00574409 -0.0193281   0.01078958]]\nJacobian from b to logits is\n[0.09748876 0.16102302 0.24190766 0.00776229]\n```\n\nExample:\n```text\ndef hessian(f):\n    return jacfwd(jacrev(f))\n\nH = hessian(f)(W)\nprint(\"hessian, with shape\", H.shape)\nprint(H)\n```\n\nExample:\n```text\nhessian, with shape (4, 3, 3)\n[[[ 0.02058932  0.04434624  0.03048803]\n  [ 0.04434623  0.09551499  0.06566654]\n  [ 0.03048803  0.06566655  0.04514575]]\n\n [[-0.0743913   0.09129842 -0.01268033]\n  [ 0.09129842 -0.11204806  0.01556223]\n  [-0.01268034  0.01556223 -0.00216142]]\n\n [[ 0.01176856  0.00135791 -0.02942139]\n  [ 0.00135791  0.00015668 -0.00339478]\n  [-0.0294214  -0.00339478  0.07355348]]\n\n [[-0.00418412  0.014079   -0.00785936]\n  [ 0.014079   -0.04737393  0.02644569]\n  [-0.00785936  0.02644569 -0.01476286]]]\n```\n\nExample:\n```text\nfrom jax import jvp\n\n# Isolate the function from the weight matrix to the predictions\nf = lambda W: predict(W, b, inputs)\n\nkey, subkey = random.split(key)\nv = random.normal(subkey, W.shape)\n\n# Push forward the vector `v` along `f` evaluated at `W`\ny, u = jvp(f, (W,), (v,))\n```\n\nExample:\n```text\njvp :: (a -> b) -> a -> T a -> (b, T b)\n```\n\nExample:\n```text\nfrom jax import vjp\n\n# Isolate the function from the weight matrix to the predictions\nf = lambda W: predict(W, b, inputs)\n\ny, vjp_fun = vjp(f, W)\n\nkey, subkey = random.split(key)\nu = random.normal(subkey, y.shape)\n\n# Pull back the covector `u` along `f` evaluated at `W`\nv = vjp_fun(u)\n```\n\nExample:\n```text\nvjp :: (a -> b) -> a -> (b, CT b -> CT a)\n```\n\nExample:\n```text\nfrom jax import vjp\n\ndef vgrad(f, x):\n  y, vjp_fn = vjp(f, x)\n  return vjp_fn(jnp.ones(y.shape))[0]\n\nprint(vgrad(lambda x: 3*x**2, jnp.ones((2, 2))))\n```\n\nExample:\n```text\n[[6. 6.]\n [6. 6.]]\n```\n\nExample:\n```text\nfrom jax import jvp, grad\n\n# forward-over-reverse\ndef hvp(f, primals, tangents):\n  return jvp(grad(f), primals, tangents)[1]\n```\n\nExample:\n```text\ndef f(X):\n  return jnp.sum(jnp.tanh(X)**2)\n\nkey, subkey1, subkey2 = random.split(key, 3)\nX = random.normal(subkey1, (30, 40))\nV = random.normal(subkey2, (30, 40))\n\nans1 = hvp(f, (X,), (V,))\nans2 = jnp.tensordot(hessian(f)(X), V, 2)\n\nprint(jnp.allclose(ans1, ans2, 1e-4, 1e-4))\n```\n\nExample:\n```text\nTrue\n```\n\nExample:\n```text\n# reverse-over-forward\ndef hvp_revfwd(f, primals, tangents):\n  g = lambda primals: jvp(f, primals, tangents)[1]\n  return grad(g)(primals)\n```\n\nExample:\n```text\n# reverse-over-reverse, only works for single arguments\ndef hvp_revrev(f, primals, tangents):\n  x, = primals\n  v, = tangents\n  return grad(lambda x: jnp.vdot(grad(f)(x), v))(x)\n\n\nprint(\"Forward over reverse\")\n%timeit -n10 -r3 hvp(f, (X,), (V,))\nprint(\"Reverse over forward\")\n%timeit -n10 -r3 hvp_revfwd(f, (X,), (V,))\nprint(\"Reverse over reverse\")\n%timeit -n10 -r3 hvp_revrev(f, (X,), (V,))\n\nprint(\"Naive full Hessian materialization\")\n%timeit -n10 -r3 jnp.tensordot(hessian(f)(X), V, 2)\n```\n\nExample:\n```text\nForward over reverse\n2.69 ms ± 110 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)\nReverse over forward\nThe slowest run took 5.09 times longer than the fastest. This could mean that an intermediate result is being cached.\n9.27 ms ± 7.55 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)\nReverse over reverse\n11.5 ms ± 7.34 ms per loop (mean ± std. dev. of 3 runs, 10 loops each)\nNaive full Hessian materialization\n41.8 ms ± 734 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)\n```\n\nExample:\n```text\n# Isolate the function from the weight matrix to the predictions\nf = lambda W: predict(W, b, inputs)\n\n# Pull back the covectors `m_i` along `f`, evaluated at `W`, for all `i`.\n# First, use a list comprehension to loop over rows in the matrix M.\ndef loop_mjp(f, x, M):\n    y, vjp_fun = vjp(f, x)\n    return jnp.vstack([jnp.asarray(vjp_fun(mi)) for mi in M])\n\n# Now, use vmap to build a computation that does a single fast matrix-matrix\n# multiply, rather than an outer loop over vector-matrix multiplies.\ndef vmap_mjp(f, x, M):\n    y, vjp_fun = vjp(f, x)\n    outs, = vmap(vjp_fun)(M)\n    return outs\n\nkey = random.key(0)\nnum_covecs = 128\nU = random.normal(key, (num_covecs,) + y.shape)\n\nloop_vs = loop_mjp(f, W, M=U)\nprint('Non-vmapped Matrix-Jacobian product')\n%timeit -n10 -r3 loop_mjp(f, W, M=U)\n\nprint('\\nVmapped Matrix-Jacobian product')\nvmap_vs = vmap_mjp(f, W, M=U)\n%timeit -n10 -r3 vmap_mjp(f, W, M=U)\n\nassert jnp.allclose(loop_vs, vmap_vs), 'Vmap and non-vmapped Matrix-Jacobian Products should be identical'\n```\n\nExample:\n```text\nNon-vmapped Matrix-Jacobian product\n59.4 ms ± 249 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)\n\nVmapped Matrix-Jacobian product\n3.38 ms ± 36 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)\n```\n\nExample:\n```text\ndef loop_jmp(f, W, M):\n    # jvp immediately returns the primal and tangent values as a tuple,\n    # so we'll compute and select the tangents in a list comprehension\n    return jnp.vstack([jvp(f, (W,), (mi,))[1] for mi in M])\n\ndef vmap_jmp(f, W, M):\n    _jvp = lambda s: jvp(f, (W,), (s,))[1]\n    return vmap(_jvp)(M)\n\nnum_vecs = 128\nS = random.normal(key, (num_vecs,) + W.shape)\n\nloop_vs = loop_jmp(f, W, M=S)\nprint('Non-vmapped Jacobian-Matrix product')\n%timeit -n10 -r3 loop_jmp(f, W, M=S)\nvmap_vs = vmap_jmp(f, W, M=S)\nprint('\\nVmapped Jacobian-Matrix product')\n%timeit -n10 -r3 vmap_jmp(f, W, M=S)\n\nassert jnp.allclose(loop_vs, vmap_vs), 'Vmap and non-vmapped Jacobian-Matrix products should be identical'\n```\n\nExample:\n```text\nNon-vmapped Jacobian-Matrix product\n79.2 ms ± 23 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)\n\nVmapped Jacobian-Matrix product\n1.15 ms ± 40.3 μs per loop (mean ± std. dev. of 3 runs, 10 loops each)\n```\n\nExample:\n```text\nfrom jax import jacrev as builtin_jacrev\n\ndef our_jacrev(f):\n    def jacfun(x):\n        y, vjp_fun = vjp(f, x)\n        # Use vmap to do a matrix-Jacobian product.\n        # Here, the matrix is the Euclidean basis, so we get all\n        # entries in the Jacobian at once.\n        J, = vmap(vjp_fun, in_axes=0)(jnp.eye(len(y)))\n        return J\n    return jacfun\n\nassert jnp.allclose(builtin_jacrev(f)(W), our_jacrev(f)(W)), 'Incorrect reverse-mode Jacobian results!'\n```\n\nExample:\n```text\nfrom jax import jacfwd as builtin_jacfwd\n\ndef our_jacfwd(f):\n    def jacfun(x):\n        _jvp = lambda s: jvp(f, (x,), (s,))[1]\n        Jt = vmap(_jvp, in_axes=1)(jnp.eye(len(x)))\n        return jnp.transpose(Jt)\n    return jacfun\n\nassert jnp.allclose(builtin_jacfwd(f)(W), our_jacfwd(f)(W)), 'Incorrect forward-mode Jacobian results!'\n```\n\nExample:\n```text\ndef f(x):\n    try:\n        if x < 3:\n            return 2 * x ** 3\n        else:\n            raise ValueError\n    except ValueError:\n        return jnp.pi * x\n\ny, f_vjp = vjp(f, 4.)\nprint(jit(f_vjp)(1.))\n```\n\nExample:\n```text\n(Array(3.1415927, dtype=float32, weak_type=True),)\n```\n\nExample:\n```text\ndef f(z):\n  x, y = jnp.real(z), jnp.imag(z)\n  return u(x, y) + v(x, y) * 1j\n\ndef g(x, y):\n  return (u(x, y), v(x, y))\n```\n\nExample:\n```text\ndef check(seed):\n  key = random.key(seed)\n\n  # random coeffs for u and v\n  key, subkey = random.split(key)\n  a, b, c, d = random.uniform(subkey, (4,))\n\n  def fun(z):\n    x, y = jnp.real(z), jnp.imag(z)\n    return u(x, y) + v(x, y) * 1j\n\n  def u(x, y):\n    return a * x + b * y\n\n  def v(x, y):\n    return c * x + d * y\n\n  # primal point\n  key, subkey = random.split(key)\n  x, y = random.uniform(subkey, (2,))\n  z = x + y * 1j\n\n  # tangent vector\n  key, subkey = random.split(key)\n  c, d = random.uniform(subkey, (2,))\n  z_dot = c + d * 1j\n\n  # check jvp\n  _, ans = jvp(fun, (z,), (z_dot,))\n  expected = (grad(u, 0)(x, y) * c +\n              grad(u, 1)(x, y) * d +\n              grad(v, 0)(x, y) * c * 1j+\n              grad(v, 1)(x, y) * d * 1j)\n  print(jnp.allclose(ans, expected))\n```\n\nExample:\n```text\ncheck(0)\ncheck(1)\ncheck(2)\n```\n\nExample:\n```text\nTrue\nTrue\nTrue\n```\n\nExample:\n```text\ndef check(seed):\n  key = random.key(seed)\n\n  # random coeffs for u and v\n  key, subkey = random.split(key)\n  a, b, c, d = random.uniform(subkey, (4,))\n\n  def fun(z):\n    x, y = jnp.real(z), jnp.imag(z)\n    return u(x, y) + v(x, y) * 1j\n\n  def u(x, y):\n    return a * x + b * y\n\n  def v(x, y):\n    return c * x + d * y\n\n  # primal point\n  key, subkey = random.split(key)\n  x, y = random.uniform(subkey, (2,))\n  z = x + y * 1j\n\n  # cotangent vector\n  key, subkey = random.split(key)\n  c, d = random.uniform(subkey, (2,))\n  z_bar = jnp.array(c + d * 1j)  # for dtype control\n\n  # check vjp\n  _, fun_vjp = vjp(fun, z)\n  ans, = fun_vjp(z_bar)\n  expected = (grad(u, 0)(x, y) * c +\n              grad(v, 0)(x, y) * (-d) +\n              grad(u, 1)(x, y) * c * (-1j) +\n              grad(v, 1)(x, y) * (-d) * (-1j))\n  assert jnp.allclose(ans, expected, atol=1e-5, rtol=1e-5)\n```\n\nExample:\n```text\ndef f(z):\n  x, y = jnp.real(z), jnp.imag(z)\n  return x**2 + y**2\n\nz = 3. + 4j\ngrad(f)(z)\n```\n\nExample:\n```text\nArray(6.-8.j, dtype=complex64)\n```\n\nExample:\n```text\ndef f(z):\n  return jnp.sin(z)\n\nz = 3. + 4j\ngrad(f, holomorphic=True)(z)\n```\n\nExample:\n```text\nArray(-27.034946-3.8511534j, dtype=complex64, weak_type=True)\n```\n\nExample:\n```text\ndef f(z):\n  return jnp.conjugate(z)\n\nz = 3. + 4j\ngrad(f, holomorphic=True)(z)  # f is not actually holomorphic!\n```\n\nExample:\n```text\nArray(1.-0.j, dtype=complex64, weak_type=True)\n```\n\nExample:\n```text\nA = jnp.array([[5.,    2.+3j,    5j],\n              [2.-3j,   7.,  1.+7j],\n              [-5j,  1.-7j,    12.]])\n\ndef f(X):\n    L = jnp.linalg.cholesky(X)\n    return jnp.sum((L - jnp.sin(L))**2)\n\ngrad(f, holomorphic=True)(A)\n```\n\nExample:\n```text\nArray([[-0.7534186  +0.j       , -3.0509028 -10.940544j ,\n         5.9896846  +3.5423026j],\n       [-3.0509028 +10.940544j , -8.904491   +0.j       ,\n        -5.1351523  -6.559373j ],\n       [ 5.9896846  -3.5423026j, -5.1351523  +6.559373j ,\n         0.01320427 +0.j       ]], dtype=complex64)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.772Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":55,"totalLines":660,"estimatedTokens":3612}}75{"id":"doc-the_training_cookbook_jax_documentation-cb60368e","source":"documentation","title":"The Training Cookbook — JAX documentation","url":"https://docs.jax.dev/en/latest/the-training-cookbook.html","text":"Example:\n```text\ndef train_loop(config: Config):\n  record_writer = RecordWriter()\n  train_state = init_train_state(config)\n  train_state = jax.tree.map(jax.ref.new_ref, train_state)\n  batch = iter(get_dataset_on_device(config))\n  for step in range(config.num_train_steps):\n    metrics = train_step(config, train_state, next(batch))\n    record_writer({\"step\": step} | metrics)\n```\n\nExample:\n```text\nparam = jnp.zeros((256, 192), out_sharding=jax.P(\"gpu\", None))\n```\n\nExample:\n```text\n@jax.jit\ndef init_train_state(config: Config) -> dot_dict:\n  train_state = dot_dict()\n  train_state.params = init_param_state(config)\n  train_state.opt = jax.tree.map(init_adam_state, train_state.params)\n  return train_state\n```\n\nExample:\n```text\ndef init_param_state(config: Config) -> dot_dict:\n  root_key = jax.random.key(config.param_seed)\n  key = map(ft.partial(jax.random.fold_in, root_key), it.count())\n  zero_init = jax.nn.initializers.constant(0.0)\n  he_init = jax.nn.initializers.he_normal(1, 1)\n  dtype = config.dtype\n\n  params = dot_dict(\n    pos_embed=zero_init(next(key), (config.seq_length, config.embed_dim), dtype, config.pos_embed),\n    layers=dot_dict(),\n  )\n  params.embedding = he_init(next(key), (config.vocab_size, config.embed_dim), dtype, config.embed)\n  params.linear_in = dot_dict(\n    kernel=he_init(next(key), (1, config.embed_dim), dtype, config.in_kernel),\n    bias=zero_init(next(key), (config.embed_dim,), dtype, config.in_bias),\n  )\n  params.linear_out = dot_dict(\n    kernel=he_init(next(key), (config.embed_dim, config.vocab_size), dtype, config.out_kernel),\n  )\n  for layer in range(config.num_layers):\n    qkv_shape = (3, config.embed_dim, config.num_heads, config.head_dim)\n    out_shape = (config.num_heads, config.head_dim, config.embed_dim)\n    params.layers[layer] = dot_dict(\n      attention=dot_dict(\n        qkv=he_init(next(key), qkv_shape, dtype, config.att_qkv),\n        out=he_init(next(key), out_shape, dtype, config.att_out),\n      ),\n      mlp=dot_dict(\n        in_kernel=he_init(next(key), (config.embed_dim, config.mlp_dim), dtype, config.mlp_in),\n        out_kernel=he_init(next(key), (config.mlp_dim, config.embed_dim), dtype, config.mlp_out),\n      ),\n    )\n  return params\n```\n\nExample:\n```text\nclass Initializer(Protocol):\n    def __call__(self, key, shape, dtype, out_sharding) -> jax.Array:\n        ...\n```\n\nExample:\n```text\ndef init_adam_state(param: jax.Array) -> dot_dict:\n  adam_state = dot_dict(mu=jnp.zeros_like(param), nu=jnp.zeros_like(param), count=jnp.array(0))\n  return adam_state\n```\n\nExample:\n```text\n@jax.jit\ndef train_step(config: Config, train_state: dot_dict, batch: dict) -> dict:\n  def loss_fn(params):\n    logits = model_apply(config, params, batch[\"observed_ids\"])\n    labels = jax.nn.one_hot(batch[\"target_ids\"], config.vocab_size)\n    return -(labels * jax.nn.log_softmax(logits)).mean()\n\n  params = jax.tree.map(jax.ref.get, train_state.params)\n  loss, grad = jax.value_and_grad(loss_fn)(params)\n  jax.tree.map(ft.partial(adam_update, config), train_state.params, grad, train_state.opt)\n  metrics = {\"train_loss\": loss}\n  return metrics\n```\n\nExample:\n```text\ndef model_apply(config: Config, params: dot_dict, tokens: jax.Array) -> jax.Array:\n  out = params.embedding.at[tokens].get(out_sharding=config.act_seq)\n  out += params.pos_embed\n  del tokens\n\n  for layer in range(config.num_layers):\n    block = params.layers[layer]\n    att_skip = out  # 1 billion dollars in venture capital funding please\n    qkv = jnp.einsum(\"bsd,3dkh->bs3kh\", out, block.attention.qkv, out_sharding=config.act_att)\n    out = jax.nn.dot_product_attention(qkv[:, :, 0, :], qkv[:, :, 1, :], qkv[:, :, 2, :], is_causal=True)\n    out = jnp.einsum(\"bskh,khd->bsd\", out, block.attention.out, out_sharding=config.act_seq)\n    out += att_skip\n    out *= jax.lax.rsqrt(jnp.linalg.norm(out, axis=-1, keepdims=True) + 1e-6)\n\n    mlp_skip = out  # machine learning circa 1986\n    out = jnp.einsum(\"bsd,dh->bsh\", out, block.mlp.in_kernel, out_sharding=config.act_hidden)\n    out = jax.nn.gelu(out)\n    out = jnp.einsum(\"bsh,hd->bsd\", out, block.mlp.out_kernel, out_sharding=config.act_seq)\n    out += mlp_skip\n    out *= jax.lax.rsqrt(jnp.linalg.norm(out, axis=-1, keepdims=True) + 1e-6)\n\n  logits = jnp.einsum(\"bsd,dl->bsl\", out, params.linear_out.kernel, out_sharding=config.act_seq)\n  return logits\n```\n\nExample:\n```text\nloss, grad = jax.value_and_grad(loss_fn)(params)\n```\n\nExample:\n```text\njax.tree.map(ft.partial(adam_update, config), train_state.params, grad, train_state.opt)\n```\n\nExample:\n```text\ndef adam_update(config: Config, param: jax.Ref, grad: jax.Array, adam_state: dot_dict):\n  adam_state.mu[...] = (1 - config.beta_1) * adam_state.mu[...] + config.beta_1 * grad\n  adam_state.nu[...] = (1 - config.beta_2) * adam_state.nu[...] + config.beta_2 * grad**2\n  adam_state.count[...] += 1\n\n  mu_hat = adam_state.mu[...] / (1 - config.beta_1 ** adam_state.count[...])\n  nu_hat = adam_state.nu[...] / (1 - config.beta_2 ** adam_state.count[...])\n  param[...] -= config.learning_rate * mu_hat / (jnp.sqrt(nu_hat + config.eps_root) + config.eps)\n```\n\nExample:\n```text\nfor step in range(config.num_train_steps):\n  metrics = train_step(config, train_state, next(batch))\n```\n\nExample:\n```text\n---\ndisplayMode: compact\n---\ngantt\n    title Synchronous Dispatch: No Overlap\n    axisFormat %\n\n    section Host\n    next(batch) :gb0, 0, 1000s\n    next(batch) :gb1, after ajc0, 1000s\n    next(batch) :gb2, after ajc1, 1000s\n\n    section Accelerator\n\n    train_step 0 :ajc0, after gb0, 2000s\n    train_step 1 :ajc1, after gb1, 2000s\n```\n\nExample:\n```text\n---\ndisplayMode: compact\n---\ngantt\n    title JAX Asynchronous Dispatch: Host-Device Overlap\n    axisFormat %\n\n    section Host\n    %% Task: id, name, start, duration_or_end\n    next(batch) :gb0, 0, 1000s\n    next(batch) :gb1, after gb0, 1000s\n    next(batch) :gb2, after gb1, 1000s\n    next(batch) :gb3, after jc0, 1000s\n    next(batch) :gb4, after jc1, 1000s\n\n    section Accelerator\n    %% Task: id, name, start, duration_or_end\n    train_step 0 :jc0, after gb1, 2000s\n    train_step 1 :jc1, after jc0, 2000s\n    train_step 2 :jc2, after jc1, 2000s\n```\n\nExample:\n```text\nmetrics = train_step(config, train_state, next(batch))\nprint({\"step\": step} | metrics)\n```\n\nExample:\n```text\nclass RecordWriter:\n  prev_metrics = None\n\n  def __call__(self, cur_metrics: dict):\n    self.prev_metrics, log_metrics = cur_metrics, self.prev_metrics\n    if log_metrics is None:\n      return\n    print(*it.starmap(\"{}: {}\".format, log_metrics.items()), sep=\"\\t\")\n```\n\nExample:\n```text\nmetrics = train_step(config, train_state, next(batch))\n```\n\nExample:\n```text\ndef learning_rate(count, init_value: float = 1e-4, decay_steps: int = 10_000, alpha: float = 1e-6):\n    cosine_decay = 0.5 * (1 + jnp.cos(jnp.pi * jnp.minimum(count, decay_steps) / decay_steps))\n    return init_value * (1 - alpha) * cosine_decay\n```\n\nExample:\n```text\nmetrics = train_step(config, train_state, next(batch))\nrecord_writer({\"step\": step, \"learning_rate\": learning_rate(step)} | metrics)\n```\n\nExample:\n```text\nmetrics = train_step(config, train_state, next(batch))\nwith jax.default_device('cpu'):\n  record_writer({\"step\": step, \"learning_rate\": learning_rate(step)} | metrics)\n```\n\nExample:\n```text\ndef get_dataset_on_device(config: Config) -> Iterator[dict[str, jax.Array]]:\n  datset = get_dataset(config)\n  sharding = jax.P(config.mesh_axis_names)\n  return map(ft.partial(jax.make_array_from_process_local_data, sharding), datset)\n```\n\nExample:\n```text\n@jax.tree_util.register_static\n@dataclass(kw_only=True, frozen=True)\nclass Config:\n  mesh_axis_names: tuple[str, ...] = (\"fsdp\",)\n  mesh_shape: tuple[int, ...] = (8,)\n  seq_length: int = 128\n\n  num_train_steps: int = 10**6\n  host_batch_size: int = 16\n  learning_rate: float = 1e-4\n  beta_1: float = 0.9\n  beta_2: float = 0.999\n  eps: float = 1e-8\n  eps_root: float = 0.0\n\n  param_seed: int = 12738\n  num_layers: int = 4\n  embed_dim: int = 512\n  mlp_dim: int = 512 * 4\n  vocab_size: int = 2**8  # uint8 ascii encoding\n  num_heads: int = 8\n  head_dim: int = 128\n  dtype: str = \"bfloat16\"\n\n  embed: jax.P = jax.P(None, None)\n  pos_embed: jax.P = jax.P(None, None)\n  att_qkv: jax.P = jax.P(None, \"fsdp\", None, None)\n  att_out: jax.P = jax.P(\"fsdp\", None, None)\n  mlp_in: jax.P = jax.P(\"fsdp\", None)\n  mlp_out: jax.P = jax.P(None, \"fsdp\")\n  in_kernel: jax.P = jax.P(None, None)\n  in_bias: jax.P = jax.P(None)\n  out_kernel: jax.P = jax.P(\"fsdp\", None)\n  out_bias: jax.P = jax.P(None)\n\n  act_ids: jax.P = jax.P(\"fsdp\")\n  act_seq: jax.P = jax.P(\"fsdp\", None, None)\n  act_att: jax.P = jax.P(\"fsdp\", None, None, None)\n  act_hidden: jax.P = jax.P(\"fsdp\", None, None)\n\n  def __post_init__(self):\n    mesh = jax.make_mesh(self.mesh_shape, self.mesh_axis_names, len(self.mesh_shape) * (AxisType.Explicit,))\n    jax.sharding.set_mesh(mesh)\n```\n\nExample:\n```text\nmesh = jax.sharding.Mesh(jax.devices(), ('devices',))\n```\n\nExample:\n```text\npos_embed = jax.P(None, None)\natt_qkv = jax.P(None, None, None, None)\natt_out = jax.P(None, None, None)\nmlp_in = jax.P(None, None)\nmlp_out = jax.P(None, None)\nin_kernel = jax.P(None, None)\nin_bias = jax.P(None)\nout_kernel = jax.P(None, None)\nout_bias = jax.P(None)\n```\n\nExample:\n```text\nact_ids = jax.P(\"devices\")\nact_seq = jax.P(\"devices\", None, None)\nact_att = jax.P(\"devices\", None, None, None)\nact_hidden = jax.P(\"devices\", None, None)\n```\n\nExample:\n```text\nmesh = jax.make_mesh((128*4,), (\"fsdp\",))\n```\n\nExample:\n```text\npos_embed = jax.P(None, None)\natt_qkv = jax.P(None, \"fsdp\", None, None)\natt_out = jax.P(\"fsdp\", None, None)\nmlp_in = jax.P(\"fsdp\", None)\nmlp_out = jax.P(None, \"fsdp\")\nin_kernel = jax.P(None, None)\nin_bias = jax.P(None)\nout_kernel = jax.P(\"fsdp\", None)\nout_bias = jax.P(None)\n```\n\nExample:\n```text\nact_ids = jax.P(\"fsdp\")\nact_seq = jax.P(\"fsdp\", None, None)\nact_att = jax.P(\"fsdp\", None, None, None)\nact_hidden = jax.P(\"fsdp\", None, None)\n```\n\nExample:\n```text\nmesh = jax.make_mesh((128,4), (\"fsdp\", \"tensor\"))\n```\n\nExample:\n```text\npos_embed = jax.P(None, \"tensor\")\natt_qkv = jax.P(None, \"fsdp\", \"tensor\", None)\natt_out = jax.P(\"fsdp\", None, None)\nmlp_in = jax.P(\"fsdp\", \"tensor\")\nmlp_out = jax.P(\"tensor\", \"fsdp\")\nin_kernel = jax.P(None, None)\nin_bias = jax.P(None)\nout_kernel = jax.P(\"fsdp\", None)\nout_bias = jax.P(None)\n```\n\nExample:\n```text\nact_ids = jax.P(\"fsdp\")\nact_seq = jax.P(\"fsdp\", None, None)\nact_att = jax.P(\"fsdp\", None, \"tensor\", None)\nact_hidden = jax.P(\"fsdp\", None, \"tensor\")\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.777Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":31,"totalLines":370,"estimatedTokens":2615}}76{"id":"doc-custom_derivative_rules_with_hijax_primitives_ja-6d3c0ad5","source":"documentation","title":"Custom derivative rules with hijax primitives — JAX documentation","url":"https://docs.jax.dev/en/latest/hijax_custom_derivatives.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.experimental.hijax import VJPHiPrimitive\n\nclass SinTimesY(VJPHiPrimitive):\n  def __init__(self, x_aval, y_aval):\n    self.in_avals = (x_aval, y_aval)  # input types\n    self.out_aval = x_aval            # output type\n    self.params = {}                  # static parameters (none here)\n    super().__init__()\n\n  # Implementation, used for evaluation and lowering (e.g. under jit).\n  def expand(self, x, y):\n    return jnp.sin(x) * y\n\n  # Reverse-mode: forward pass returns (primal_out, residuals).\n  def vjp_fwd(self, nzs_in, x, y):\n    return self(x, y), (jnp.cos(x), jnp.sin(x), y)\n\n  # Reverse-mode: backward pass maps (residuals, output cotangent) to a tuple\n  # of input cotangents.\n  def vjp_bwd_retval(self, res, g):\n    cos_x, sin_x, y = res\n    return (cos_x * g * y, sin_x * g)\n\n  # Forward-mode rule (optional, only needed under e.g. jax.jvp).\n  def jvp(self, primals, tangents):\n    (x, y), (x_dot, y_dot) = primals, tangents\n    return self(x, y), jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot\n\ndef f(x, y):\n  return SinTimesY(jax.typeof(x), jax.typeof(y))(x, y)\n```\n\nExample:\n```text\nfrom jax import jvp, grad\n\nprint(f(2., 3.))\ny, y_dot = jvp(f, (2., 3.), (1., 0.))\nprint(y)\nprint(y_dot)\nprint(grad(f)(2., 3.))\n```\n\nExample:\n```text\n2.7278922\n2.7278922\n-1.2484405\n-1.2484405\n```\n\nExample:\n```text\ndef log1pexp(x):\n  return jnp.log(1. + jnp.exp(x))\n\nlog1pexp(3.)\n```\n\nExample:\n```text\nArray(3.0485873, dtype=float32, weak_type=True)\n```\n\nExample:\n```text\nfrom jax import jit, grad, vmap\n\nprint(jit(log1pexp)(3.))\nprint(jit(grad(log1pexp))(3.))\nprint(vmap(jit(grad(log1pexp)))(jnp.arange(3.)))\n```\n\nExample:\n```text\n3.0485873\n0.95257413\n[0.5       0.7310586 0.8807971]\n```\n\nExample:\n```text\nprint(grad(log1pexp)(100.))\n```\n\nExample:\n```text\nnan\n```\n\nExample:\n```text\njit(grad(log1pexp)).trace(100.).jaxpr\n```\n\nExample:\n```text\n{ lambda ; a:f32[]. let\n    b:f32[] = exp a\n    c:f32[] = add 1.0:f32[] b\n    _:f32[] = log c\n    d:f32[] = div 1.0:f32[] c\n    e:f32[] = mul d b\n  in (e,) }\n```\n\nExample:\n```text\nclass Log1pExp(VJPHiPrimitive):\n  def __init__(self, x_aval):\n    self.in_avals = (x_aval,)\n    self.out_aval = x_aval\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x):\n    return jnp.log(1. + jnp.exp(x))\n\n  def vjp_fwd(self, nzs_in, x):\n    return self(x), x\n\n  def vjp_bwd_retval(self, x, g):\n    return ((1 - 1/(1 + jnp.exp(x))) * g,)\n\n  def jvp(self, primals, tangents):\n    (x,), (x_dot,) = primals, tangents\n    return self(x), (1 - 1/(1 + jnp.exp(x))) * x_dot\n\n  def batch_dim_rule(self, axis_data, in_dims):\n    return in_dims[0]\n\ndef log1pexp(x):\n  return Log1pExp(jax.typeof(x))(x)\n```\n\nExample:\n```text\n1.0\n```\n\nExample:\n```text\nprint(jit(log1pexp)(3.))\nprint(jit(grad(log1pexp))(3.))\nprint(vmap(jit(grad(log1pexp)))(jnp.arange(3.)))\n```\n\nExample:\n```text\n{ lambda ; a:f32[]. let\n    _:f32[] = call_hi_primitive[_prim=Log1pExp[{}]] a\n    b:f32[] = exp a\n    c:f32[] = add 1.0:f32[] b\n    d:f32[] = div 1.0:f32[] c\n    e:f32[] = sub 1.0:f32[] d\n    f:f32[] = mul e 1.0:f32[]\n  in (f,) }\n```\n\nExample:\n```text\ndef f(x):\n  return x / (1 + jnp.sqrt(x))\n```\n\nExample:\n```text\nprint(grad(f)(0.))\n```\n\nExample:\n```text\nclass FOnRPlus(VJPHiPrimitive):\n  def __init__(self, x_aval):\n    self.in_avals = (x_aval,)\n    self.out_aval = x_aval\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x):\n    return x / (1 + jnp.sqrt(x))\n\n  def vjp_fwd(self, nzs_in, x):\n    return self(x), x\n\n  def vjp_bwd_retval(self, x, g):\n    return (((jnp.sqrt(x) + 2) / (2 * (jnp.sqrt(x) + 1)**2)) * g,)\n\ndef f(x):\n  return FOnRPlus(jax.typeof(x))(x)\n```\n\nExample:\n```text\nclass ClipGradient(VJPHiPrimitive):\n  def __init__(self, lo_aval, hi_aval, x_aval):\n    self.in_avals = (lo_aval, hi_aval, x_aval)\n    self.out_aval = x_aval\n    self.params = {}\n    super().__init__()\n\n  def expand(self, lo, hi, x):\n    return x  # identity function\n\n  def vjp_fwd(self, nzs_in, lo, hi, x):\n    return self(lo, hi, x), (lo, hi)  # save bounds as residuals\n\n  def vjp_bwd_retval(self, res, g):\n    lo, hi = res\n    return (None, None, jnp.clip(g, lo, hi))  # None: zero cotangents for lo, hi\n\n  def batch_dim_rule(self, axis_data, in_dims):\n    return in_dims[2]\n\ndef clip_gradient(lo, hi, x):\n  return ClipGradient(jax.typeof(lo), jax.typeof(hi), jax.typeof(x))(lo, hi, x)\n```\n\nExample:\n```text\nimport matplotlib.pyplot as plt\n\nt = jnp.linspace(0, 10, 1000)\n\nplt.plot(jnp.sin(t))\nplt.plot(vmap(grad(jnp.sin))(t))\n```\n\nExample:\n```text\n[<matplotlib.lines.Line2D at 0x77e52d2a2a20>]\n```\n\nExample:\n```text\ndef clip_sin(x):\n  x = clip_gradient(-0.75, 0.75, x)\n  return jnp.sin(x)\n\nplt.plot(clip_sin(t))\nplt.plot(vmap(grad(clip_sin))(t))\n```\n\nExample:\n```text\n[<matplotlib.lines.Line2D at 0x77e52cf63e30>]\n```\n\nExample:\n```text\nimport pdb\n\nclass Debug(VJPHiPrimitive):\n  def __init__(self, x_aval):\n    self.in_avals = (x_aval,)\n    self.out_aval = x_aval\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x):\n    return x  # acts like identity\n\n  def vjp_fwd(self, nzs_in, x):\n    return self(x), x\n\n  def vjp_bwd_retval(self, x, g):\n    pdb.set_trace()\n    return (g,)\n\ndef debug(x):\n  return Debug(jax.typeof(x))(x)\n\ndef foo(x):\n  y = x ** 2\n  y = debug(y)  # insert pdb in corresponding backward pass step\n  return jnp.sin(y)\n```\n\nExample:\n```text\njax.grad(foo)(3.)\n\n> <ipython-input-113-b19a2dc1abf7>(12)vjp_bwd_retval()\n-> return (g,)\n(Pdb) p x\nArray(9., dtype=float32)\n(Pdb) p g\nArray(-0.91113025, dtype=float32)\n(Pdb) q\n```\n\nExample:\n```text\nfrom jax.lax import while_loop\n\ndef fixed_point(f, a, x_guess):\n  def cond_fun(carry):\n    x_prev, x = carry\n    return jnp.abs(x_prev - x) > 1e-6\n\n  def body_fun(carry):\n    _, x = carry\n    return x, f(a, x)\n\n  _, x_star = while_loop(cond_fun, body_fun, (x_guess, f(a, x_guess)))\n  return x_star\n```\n\nExample:\n```text\ndef newton_sqrt(a):\n  update = lambda a, x: 0.5 * (x + a / x)\n  return fixed_point(update, a, a)\n```\n\nExample:\n```text\nprint(newton_sqrt(2.))\n```\n\nExample:\n```text\n1.4142135\n```\n\nExample:\n```text\nprint(jit(vmap(newton_sqrt))(jnp.array([1., 2., 3., 4.])))\n```\n\nExample:\n```text\n[1.        1.4142135 1.7320509 2.       ]\n```\n\nExample:\n```text\nfrom functools import partial\nfrom jax import vjp\n\nclass FixedPoint(VJPHiPrimitive):\n  def __init__(self, a_aval, x_aval, *, f):\n    self.in_avals = (a_aval, x_aval)\n    self.out_aval = x_aval\n    self.params = dict(f=f)\n    super().__init__()\n\n  def expand(self, a, x_guess):\n    def cond_fun(carry):\n      x_prev, x = carry\n      return jnp.abs(x_prev - x) > 1e-6\n\n    def body_fun(carry):\n      _, x = carry\n      return x, self.f(a, x)\n\n    _, x_star = while_loop(cond_fun, body_fun, (x_guess, self.f(a, x_guess)))\n    return x_star\n\n  def vjp_fwd(self, nzs_in, a, x_guess):\n    x_star = self(a, x_guess)\n    return x_star, (a, x_star)\n\n  def vjp_bwd_retval(self, res, x_star_bar):\n    a, x_star = res\n    _, vjp_a = vjp(lambda a: self.f(a, x_star), a)\n    a_bar, = vjp_a(fixed_point(partial(rev_iter, self.f),\n                               (a, x_star, x_star_bar),\n                               x_star_bar))\n    return a_bar, jnp.zeros_like(x_star)\n\ndef rev_iter(f, packed, u):\n  a, x_star, x_star_bar = packed\n  _, vjp_x = vjp(lambda x: f(a, x), x_star)\n  return x_star_bar + vjp_x(u)[0]\n\ndef fixed_point(f, a, x_guess):\n  a_aval = jax.tree.map(jax.typeof, a)\n  x_aval = jax.tree.map(jax.typeof, x_guess)\n  return FixedPoint(a_aval, x_aval, f=f)(a, x_guess)\n```\n\nExample:\n```text\nprint(grad(newton_sqrt)(2.))\nprint(grad(grad(newton_sqrt))(2.))\n```\n\nExample:\n```text\n0.35355338\n-0.088388346\n```\n\nExample:\n```text\nprint(grad(jnp.sqrt)(2.))\nprint(grad(grad(jnp.sqrt))(2.))\n```\n\nExample:\n```text\n0.35355338\n-0.08838835\n```\n\nExample:\n```text\nclass Square(VJPHiPrimitive):\n  def __init__(self, x_aval):\n    self.in_avals = (x_aval,)\n    self.out_aval = x_aval\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x):\n    return x * x\n\ndef square(x):\n  return Square(jax.typeof(x))(x)\n\nprint(square(3.))\n```\n\nExample:\n```text\n9.0\n```\n\nExample:\n```text\nvjp_fwd :: (NonZeros, a) -> (b, c)\nvjp_bwd_retval :: (c, CT b) -> CT a\n```\n\nExample:\n```text\nclass Mul(VJPHiPrimitive):\n  def __init__(self, x_aval, y_aval):\n    self.in_avals = (x_aval, y_aval)\n    self.out_aval = x_aval\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x, y):\n    return x * y\n\n  def vjp_fwd(self, nzs_in, x, y):\n    return self(x, y), (x, y)\n\n  def vjp_bwd_retval(self, res, g):\n    x, y = res\n    return (g * y, x * g)\n\ndef mul(x, y):\n  return Mul(jax.typeof(x), jax.typeof(y))(x, y)\n\nprint(grad(mul)(2., 3.))\nprint(grad(mul, 1)(2., 3.))\n```\n\nExample:\n```text\n3.0\n2.0\n```\n\nExample:\n```text\nclass Sin(VJPHiPrimitive):\n  def __init__(self, x_aval):\n    self.in_avals = (x_aval,)\n    self.out_aval = x_aval\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x):\n    return jnp.sin(x)\n\n  def jvp(self, primals, tangents):\n    (x,), (x_dot,) = primals, tangents\n    return self(x), jnp.cos(x) * x_dot\n\ndef sin(x):\n  return Sin(jax.typeof(x))(x)\n\ny, y_dot = jvp(sin, (3.,), (1.,))\nprint(y)\nprint(y_dot)\n```\n\nExample:\n```text\n0.14112\n-0.9899925\n```\n\nExample:\n```text\nfrom jax.experimental.hijax import linearize_from_jvp, vjp_from_lin\n\nclass SinAD(Sin):\n  lin, linearized = linearize_from_jvp\n  vjp_fwd, vjp_bwd_retval = vjp_from_lin\n\ndef sin(x):\n  return SinAD(jax.typeof(x))(x)\n\nprint(grad(sin)(3.))\ny, sin_lin = jax.linearize(sin, 3.)\nprint(sin_lin(1.))\nprint(grad(grad(sin))(3.))\n```\n\nExample:\n```text\n-0.9899925\n-0.9899925\n-0.14112\n```\n\nExample:\n```text\njit(mul).trace(2., 3.).jaxpr\n```\n\nExample:\n```text\n{ lambda ; a:f32[] b:f32[]. let\n    c:f32[] = call_hi_primitive[_prim=Mul[{}]] a b\n  in (c,) }\n```\n\nExample:\n```text\nclass Noisy(VJPHiPrimitive):\n  def __init__(self, x_aval):\n    self.in_avals = (x_aval,)\n    self.out_aval = x_aval\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x):\n    print('called expand!')\n    return jnp.sin(x)\n\ndef noisy(x):\n  return Noisy(jax.typeof(x))(x)\n\nprint(noisy(3.))\n```\n\nExample:\n```text\ncalled expand!\n0.14112\n```\n\nExample:\n```text\nprint(jit(noisy)(3.))\nprint(jit(noisy)(3.))  # tracing is cached: no more 'called expand!'\n```\n\nExample:\n```text\ncalled expand!\n0.14112\n0.14112\n```\n\nExample:\n```text\nclass G(VJPHiPrimitive):\n  def __init__(self, x_aval):\n    self.in_avals = (x_aval,)\n    self.out_aval = x_aval\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x):\n    if x > 0:\n      return jnp.sin(x)\n    else:\n      return jnp.cos(x)\n\n  def vjp_fwd(self, nzs_in, x):\n    return self(x), x\n\n  def vjp_bwd_retval(self, x, g):\n    if x > 0:\n      return (2 * g,)\n    else:\n      return (3 * g,)\n\ndef g(x):\n  return G(jax.typeof(x))(x)\n\nprint(grad(g)(1.))\nprint(grad(g)(-1.))\n```\n\nExample:\n```text\n2.0\n3.0\n```\n\nExample:\n```text\nclass MulV(Mul):\n  def batch_dim_rule(self, axis_data, in_dims):\n    x_dim, y_dim = in_dims\n    return y_dim if x_dim is None else x_dim\n\ndef mul(x, y):\n  return MulV(jax.typeof(x), jax.typeof(y))(x, y)\n\nx = jnp.arange(3.)\ny = jnp.arange(3.) + 1.\nprint(vmap(mul)(x, y))\nprint(vmap(mul, in_axes=(0, None))(x, 2.))\nprint(vmap(grad(mul))(x, y))\n```\n\nExample:\n```text\n[0. 2. 6.]\n[0. 2. 4.]\n[1. 2. 3.]\n```\n\nExample:\n```text\nfrom collections import namedtuple\nPoint = namedtuple(\"Point\", [\"x\", \"y\"])\n\nfrom jax.experimental.hijax import Zero, instantiate_zeros\n\nclass FPt(VJPHiPrimitive):\n  def __init__(self, pt_aval):\n    self.in_avals = (pt_aval,)\n    self.out_aval = {'a': pt_aval.x, 'b': (pt_aval.x, pt_aval.y)}\n    self.params = {}\n    super().__init__()\n\n  def expand(self, pt):\n    return {'a': pt.x ** 2, 'b': (jnp.sin(pt.x), jnp.cos(pt.y))}\n\n  def vjp_fwd(self, nzs_in, pt):\n    return self(pt), pt\n\n  def vjp_bwd_retval(self, pt, g):\n    g = jax.tree.map(instantiate_zeros, g,\n                     is_leaf=lambda x: isinstance(x, Zero))\n    a_bar, (b0_bar, b1_bar) = g['a'], g['b']\n    x_bar = 2 * pt.x * a_bar + jnp.cos(pt.x) * b0_bar\n    y_bar = -jnp.sin(pt.y) * b1_bar\n    return (Point(x_bar, y_bar),)\n\ndef f(pt):\n  return FPt(jax.tree.map(jax.typeof, pt))(pt)\n\ndef fun(pt):\n  dct = f(pt)\n  return dct['a'] + dct['b'][0]\n\npt = Point(1., 2.)\nprint(f(pt))\nprint(grad(fun)(pt))\n```\n\nExample:\n```text\n{'a': 1.0, 'b': (Array(0.84147096, dtype=float32, weak_type=True), Array(-0.41614684, dtype=float32, weak_type=True))}\nPoint(x=Array(2.5403023, dtype=float32, weak_type=True), y=Array(-0., dtype=float32, weak_type=True))\n```\n\nExample:\n```text\nclass Mul2(VJPHiPrimitive):\n  def __init__(self, x_aval, y_aval):\n    self.in_avals = (x_aval, y_aval)\n    self.out_aval = x_aval\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x, y):\n    return x * y\n\n  def vjp_fwd(self, nzs_in, x, y):\n    x_nz, y_nz = nzs_in\n    return self(x, y), (x if y_nz else None, y if x_nz else None)\n\n  def vjp_bwd_retval(self, res, g):\n    x, y = res\n    return (g * y if y is not None else None,\n            x * g if x is not None else None)\n\ndef mul2(x, y):\n  return Mul2(jax.typeof(x), jax.typeof(y))(x, y)\n\nprint(grad(mul2, 0)(2., 3.))  # nzs_in == (True, False), saves only y\nprint(grad(mul2, 1)(2., 3.))  # nzs_in == (False, True), saves only x\n```\n\nExample:\n```text\nclass Mul3(VJPHiPrimitive):\n  def __init__(self, x_aval, y_aval):\n    self.in_avals = (x_aval, y_aval)\n    self.out_aval = x_aval\n    self.params = {}\n    super().__init__()\n\n  def expand(self, x, y):\n    return x * y\n\n  def vjp_fwd(self, nzs_in, x, y):\n    return self(x, y), (x, y)\n\n  def vjp_bwd(self, res, g, x_acc, y_acc):\n    x, y = res\n    x_acc.accum(g * y)\n    y_acc.accum(x * g)\n\ndef mul3(x, y):\n  return Mul3(jax.typeof(x), jax.typeof(y))(x, y)\n\nprint(grad(mul3)(2., 3.))\n```\n\nExample:\n```text\n3.0\n```\n\nExample:\n```text\nclass Sin2(Sin):\n  def lin(self, nzs_in, x):\n    return self(x), jnp.cos(x)\n\n  def linearized(self, cos_x, x_dot):\n    return cos_x * x_dot\n\ndef sin2(x):\n  return Sin2(jax.typeof(x))(x)\n\ny, f_lin = jax.linearize(sin2, 3.)\nprint(y)\nprint(f_lin(1.))\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.780Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":61,"totalLines":759,"estimatedTokens":3494}}77{"id":"doc-jax_experimental_xla_metadata_module_jax_documen-3bbb99a6","source":"documentation","title":"jax.experimental.xla_metadata module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.experimental.xla_metadata.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.experimental.xla_metadata import must_fuse_call\n\n\n@jax.jit\ndef f(x):\n  y = jnp.sin(x)\n  z = must_fuse_call('1')(lambda x: jnp.square(x).sum())(x)\n  return y, z\n```\n\nExample:\n```text\n%xla_metadata_call.2 (Arg_0.1: f32[128]) -> f32[] {\n  %Arg_0.1 = f32[128]{0} parameter(0)\n  %square.1 = f32[128]{0} multiply(%Arg_0.1, %Arg_0.1)\n  %constant.1 = f32[] constant(0)\n  ROOT %reduce_sum.7 = f32[] reduce(%square.1, %constant.1),\n  dimensions={0},\n  to_apply=%region_0.1\n}\n\n\nENTRY main {\n  ...\n  %xla_metadata_call.1 = f32[] call(%x.1),\n  to_apply=%xla_metadata_call.2,\n  frontend_attributes={MUST_FUSE=\"1\"}\n  ...\n}\n```\n\nExample:\n```text\n%fused_computation (param_0.2: f32[128]) -> (f32[], f32[128]) {\n  %param_0.2 = f32[128]{0:T(128)} parameter(0)\n  %square.2 = f32[128]{0:T(128)} multiply(%param_0.2, %param_0.2),\n  frontend_attributes={MUST_FUSE=\"1\"}\n  %constant.2 = f32[]{:T(128)} constant(0),\n  frontend_attributes={MUST_FUSE=\"1\"}\n  %reduce_sum.1 = f32[]{:T(128)} reduce(%square.2, %constant.2),\n  dimensions={0}, to_apply=%region_0.1,\n  frontend_attributes={MUST_FUSE=\"1\"}\n  %sin.0 = f32[128]{0:T(128)} sine(%param_0.2)\n  ROOT %tuple = (f32[]{:T(128)}, f32[128]{0:T(128)}) tuple(%reduce_sum.1,\n  %sin.0)\n}\n\n\nENTRY main {\n  ...\n  %multiply_reduce_fusion = (f32[]{:T(128)}, f32[128]{0:T(128)})\n  fusion(%x.1), kind=kLoop, calls=%fused_computation,\n  frontend_attributes={MUST_FUSE=\"1\"}\n  ...\n}\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.781Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":62,"estimatedTokens":367}}78{"id":"doc-jax_numpy_fft_rfft_jax_documentation-8a771e78","source":"documentation","title":"jax.numpy.fft.rfft — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.rfft.html","text":"Example:\n```text\n>>> x = jnp.array([[1, 3, 5],\n...                [2, 4, 6]])\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.rfft(x)\nArray([[ 9.+0.j  , -3.+1.73j],\n       [12.+0.j  , -3.+1.73j]], dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.rfft(x, n=5)\nArray([[ 9.  +0.j  , -2.12-5.79j,  0.12+2.99j],\n       [12.  +0.j  , -1.62-7.33j,  0.62+3.36j]], dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.rfft(x, n=4, axis=0)\nArray([[ 3.+0.j,  7.+0.j, 11.+0.j],\n       [ 1.-2.j,  3.-4.j,  5.-6.j],\n       [-1.+0.j, -1.+0.j, -1.+0.j]], dtype=complex64)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.784Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":28,"estimatedTokens":177}}79{"id":"doc-jax_sharding_module_jax_documentation-07f4f72a","source":"documentation","title":"jax.sharding module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.sharding.html","text":"Example:\n```text\n>>> single_device_sharding = jax.sharding.SingleDeviceSharding(\n...     jax.devices()[0])\n```\n\nExample:\n```text\n>>> from jax.sharding import Mesh\n>>> from jax.sharding import PartitionSpec as P\n>>> mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ('x', 'y'))\n>>> spec = P('x', 'y')\n>>> named_sharding = jax.sharding.NamedSharding(mesh, spec)\n```\n\nExample:\n```text\n>>> from jax.sharding import Mesh\n>>> from jax.sharding import PartitionSpec as P, NamedSharding\n>>> import numpy as np\n...\n>>> # Declare a 2D mesh with axes `x` and `y`.\n>>> devices = np.array(jax.devices()).reshape(4, 2)\n>>> mesh = Mesh(devices, ('x', 'y'))\n>>> inp = np.arange(16).reshape(8, 2)\n>>> arr = jax.device_put(inp, NamedSharding(mesh, P('x', 'y')))\n>>> out = jax.jit(lambda x: x * 2)(arr)\n>>> assert out.sharding == NamedSharding(mesh, P('x', 'y'))\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.785Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":31,"estimatedTokens":216}}80{"id":"doc-jax_numpy_fft_rfft2_jax_documentation-ee580af6","source":"documentation","title":"jax.numpy.fft.rfft2 — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.numpy.fft.rfft2.html","text":"Example:\n```text\n>>> x = jnp.array([[[1, 3, 5],\n...                 [2, 4, 6]],\n...                [[7, 9, 11],\n...                 [8, 10, 12]]])\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.rfft2(x)\nArray([[[21.+0.j  , -6.+3.46j],\n        [-3.+0.j  ,  0.+0.j  ]],\n\n       [[57.+0.j  , -6.+3.46j],\n        [-3.+0.j  ,  0.+0.j  ]]], dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.rfft2(x, s=[2, 4])\nArray([[[21. +0.j, -8. -7.j,  7. +0.j],\n        [-3. +0.j,  0. +1.j, -1. +0.j]],\n\n       [[57. +0.j, -8.-19.j, 19. +0.j],\n        [-3. +0.j,  0. +1.j, -1. +0.j]]], dtype=complex64)\n```\n\nExample:\n```text\n>>> with jnp.printoptions(precision=2, suppress=True):\n...   jnp.fft.rfft2(x, s=[3, 5], axes=(0, 1))\nArray([[[ 18.   +0.j  ,  26.   +0.j  ,  34.   +0.j  ],\n        [ 11.09 -9.51j,  16.33-13.31j,  21.56-17.12j],\n        [ -0.09 -5.88j,   0.67 -8.23j,   1.44-10.58j]],\n\n      [[ -4.5 -12.99j,  -2.5 -16.45j,  -0.5 -19.92j],\n        [ -9.71 -6.3j , -10.05 -9.52j, -10.38-12.74j],\n        [ -4.95 +0.72j,  -5.78 -0.2j ,  -6.61 -1.12j]],\n\n      [[ -4.5 +12.99j,  -2.5 +16.45j,  -0.5 +19.92j],\n        [  3.47+10.11j,   6.43+11.42j,   9.38+12.74j],\n        [  3.19 +1.63j,   4.4  +1.38j,   5.61 +1.12j]]], dtype=complex64)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.785Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":44,"estimatedTokens":331}}81{"id":"doc-jax_random_module_jax_documentation-8cd23edf","source":"documentation","title":"jax.random module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.random.html","text":"Example:\n```text\n>>> seed = 1701\n>>> num_steps = 100\n>>> key = jax.random.key(seed)\n>>> for i in range(num_steps):\n...   key, subkey = jax.random.split(key)\n...   params = compiled_update(subkey, params, next(batches))\n```\n\nExample:\n```text\n>>> from jax import random\n>>> key = random.key(0)\n>>> key\nArray((), dtype=key<fry>) overlaying:\n[0 0]\n```\n\nExample:\n```text\n>>> random.uniform(key)\nArray(0.947667, dtype=float32)\n```\n\nExample:\n```text\n>>> key, subkey = random.split(key)\n>>> random.uniform(subkey)\nArray(0.00729382, dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.787Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":33,"estimatedTokens":139}}82{"id":"doc-jax_lax_module_jax_documentation-affa161c","source":"documentation","title":"jax.lax module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.lax.html","text":"Example:\n```text\n>>> algorithm = DotAlgorithm(\n...     lhs_precision_type=np.float16,\n...     rhs_precision_type=np.float16,\n...     accumulation_type=np.float32,\n... )\n>>> lhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)\n>>> rhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)\n>>> dot(lhs, rhs, precision=algorithm)  \narray([ 1.,  4.,  9., 16.], dtype=float16)\n```\n\nExample:\n```text\n>>> algorithm = DotAlgorithmPreset.F16_F16_F32\n>>> dot(lhs, rhs, precision=algorithm)  \narray([ 1.,  4.,  9., 16.], dtype=float16)\n```\n\nExample:\n```text\n>>> dot(lhs, rhs, precision=\"F16_F16_F32\")  \narray([ 1.,  4.,  9., 16.], dtype=float16)\n```\n\nExample:\n```text\n>>> dot(lhs, rhs, precision=\"F16_F16_F32\", preferred_element_type=np.float32)  \narray([ 1.,  4.,  9., 16.], dtype=float32)\n```\n\nExample:\n```text\n>>> lhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)\n>>> rhs = jnp.array([1.0, 2.0, 3.0, 4.0], dtype=np.float16)\n>>> algorithm = DotAlgorithmPreset.F16_F16_F32\n>>> dot(lhs, rhs, precision=algorithm)  \narray([ 1.,  4.,  9., 16.], dtype=float16)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.791Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":5,"totalLines":42,"estimatedTokens":268}}83{"id":"doc-jax_typing_module_jax_documentation-fb6ee20d","source":"documentation","title":"jax.typing module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.typing.html","text":"Example:\n```text\nimport numpy as np\nimport jax.numpy as jnp\nfrom jax import Array\nfrom jax.typing import ArrayLike\n\ndef my_function(x: ArrayLike) -> Array:\n  # Runtime type validation, Python 3.10 or newer:\n  if not isinstance(x, ArrayLike):\n    raise TypeError(f\"Expected arraylike input; got {x}\")\n  # Runtime type validation, any Python version:\n  if not (isinstance(x, (np.ndarray, Array)) or np.isscalar(x)):\n    raise TypeError(f\"Expected arraylike input; got {x}\")\n\n  # Convert input to jax.Array:\n  x_arr = jnp.asarray(x)\n\n  # ... do some computation; JAX functions will return Array types:\n  result = x_arr.sum(0) / x_arr.shape[0]\n\n  # return an Array\n  return result\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.792Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":1,"totalLines":26,"estimatedTokens":174}}84{"id":"doc-jax_export_module_jax_documentation-8fb8713b","source":"documentation","title":"jax.export module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.export.html","text":"Example:\n```text\n>>> from jax import export, sharding\n>>> # Prepare the exported object:\n>>> exp_mesh = sharding.Mesh(jax.devices(), (\"a\",))\n>>> exp = export.export(jax.jit(lambda x: jax.numpy.add(x, x),\n...                             in_shardings=sharding.NamedSharding(exp_mesh, sharding.PartitionSpec(\"a\")))\n...     )(np.arange(jax.device_count()))\n>>> exp.in_shardings_jax(exp_mesh)\n(NamedSharding(mesh=Mesh('a': 8, axis_types=(Auto,)), spec=P('a',), memory_kind=device),)\n>>> # Create a mesh for running the exported object\n>>> run_mesh = sharding.Mesh(jax.devices()[::-1], (\"a\",))\n>>> # Put the args and kwargs on the appropriate devices\n>>> run_arg = jax.device_put(np.arange(jax.device_count()),\n...     exp.in_shardings_jax(run_mesh)[0])\n>>> res = exp.call(run_arg)\n>>> res.addressable_shards\n[Shard(device=CpuDevice(id=7), index=(slice(0, 1, None),), replica_id=0, data=[0]),\n Shard(device=CpuDevice(id=6), index=(slice(1, 2, None),), replica_id=0, data=[2]),\n Shard(device=CpuDevice(id=5), index=(slice(2, 3, None),), replica_id=0, data=[4]),\n Shard(device=CpuDevice(id=4), index=(slice(3, 4, None),), replica_id=0, data=[6]),\n Shard(device=CpuDevice(id=3), index=(slice(4, 5, None),), replica_id=0, data=[8]),\n Shard(device=CpuDevice(id=2), index=(slice(5, 6, None),), replica_id=0, data=[10]),\n Shard(device=CpuDevice(id=1), index=(slice(6, 7, None),), replica_id=0, data=[12]),\n Shard(device=CpuDevice(id=0), index=(slice(7, 8, None),), replica_id=0, data=[14])]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.795Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":1,"totalLines":28,"estimatedTokens":374}}85{"id":"doc-jax_experimental_jet_module_jax_documentation-3d11dc86","source":"documentation","title":"jax.experimental.jet module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.experimental.jet.html","text":"Example:\n```text\n>>> import jax\n>>> import jax.numpy as np\n```\n\nExample:\n```text\n>>> h0, h1, h2 = 0.5**3., 3.*0.5**2., 6.*0.5\n>>> f, df, ddf = np.sin, np.cos, lambda *args: -np.sin(*args)\n```\n\nExample:\n```text\n>>> f0, (f1, f2) =  jet(f, (h0,), ((h1, h2),))\n>>> print(f0,  f(h0))\n0.12467473 0.12467473\n```\n\nExample:\n```text\n>>> print(f1, df(h0) * h1)\n0.74414825 0.74414825\n```\n\nExample:\n```text\n>>> print(f2, ddf(h0) * h1 ** 2 + df(h0) * h2)\n2.9064636 2.9064634\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.799Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":5,"totalLines":32,"estimatedTokens":120}}86{"id":"doc-jax_experimental_custom_partitioning_module_jax_-ab9d93fd","source":"documentation","title":"jax.experimental.custom_partitioning module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.experimental.custom_partitioning.html","text":"Example:\n```text\n@custom_partitioning\ndef f(*args):\n  return ...\n\ndef propagate_user_sharding(mesh, user_shape):\n  '''Update the sharding of the op from a user's shape.sharding.'''\n  user_sharding = jax.tree.map(lambda x: x.sharding, user_shape)\n\ndef partition(mesh, arg_shapes, result_shape):\n  def lower_fn(*args):\n    ... builds computation on per-device shapes ...\n  result_shardings = jax.tree.map(lambda x: x.sharding, result_shape)\n  arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)\n  # result_sharding and arg_shardings may optionally be modified and the\n  # partitioner will insert collectives to reshape.\n  return mesh, lower_fn, result_sharding, arg_shardings\n\ndef infer_sharding_from_operands(mesh, arg_shapes, shape):\n  '''Compute the result sharding from the sharding of the operands.'''\n  arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)\n\n\nf.def_partition(partition, propagate_user_sharding,\n                infer_sharding_from_operands=infer_sharding_from_operands,\n                sharding_rule='i j -> 'i j')\n```\n\nExample:\n```text\nimport jax\nfrom jax.sharding import NamedSharding\nfrom jax.experimental.custom_partitioning import custom_partitioning\nfrom jax.experimental.pjit import pjit\nfrom jax.sharding import PartitionSpec as P\nfrom jax.sharding import Mesh\nfrom jax.numpy.fft import fft\nimport regex as re\nimport numpy as np\n\n# Pattern to detect all-gather or dynamic-slice in the generated HLO\n_PATTERN = '(dynamic-slice|all-gather)'\n\n# For an N-D input, keeps sharding along the first N-1 dimensions\n# but replicate along the last dimension\ndef supported_sharding(sharding, shape):\n    rank = len(shape.shape)\n    max_shared_dims = min(len(sharding.spec), rank-1)\n    names = tuple(sharding.spec[:max_shared_dims]) + tuple(None for _ in range(rank - max_shared_dims))\n    return NamedSharding(sharding.mesh, P(*names))\n\ndef partition(mesh, arg_shapes, result_shape):\n    result_shardings = jax.tree.map(lambda x: x.sharding, result_shape)\n    arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)\n    return mesh, fft,               supported_sharding(arg_shardings[0], arg_shapes[0]),               (supported_sharding(arg_shardings[0], arg_shapes[0]),)\n\ndef infer_sharding_from_operands(mesh, arg_shapes, result_shape):\n    arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)\n    return supported_sharding(arg_shardings[0], arg_shapes[0])\n\n@custom_partitioning\ndef my_fft(x):\n    return fft(x)\n\n# Use Einsum-like notation to specify the sharding rule.\nmy_fft.def_partition(\n  infer_sharding_from_operands=infer_sharding_from_operands,\n  partition=partition,\n  sharding_rule='...i -> ...i')\n# Use SdyShardingRule object to specify the sharding rule.\nmy_fft.def_partition(\n  infer_sharding_from_operands=infer_sharding_from_operands,\n  partition=partition,\n  sharding_rule=SdyShardingRule(operand_mappings=((BATCHING, 'i'),), result_mappings=((BATCHING, 'i'),))))\n```\n\nExample:\n```text\nwith Mesh(np.array(jax.devices()), ('x',)):\n  x = np.asarray(np.random.randn(32*1024, 1024), dtype=np.complex64)\n  y = pjit(lambda x: x, in_shardings=None, out_shardings=P('x'))(x)\n  pjit_my_fft = pjit(my_fft, in_shardings=P('x'), out_shardings=P('x'))\n  pjit_fft    = pjit(fft,    in_shardings=P('x'), out_shardings=P('x'))\n  print(pjit_my_fft(y))\n  print(pjit_fft(y))\n  # dynamic-slice or all-gather are not present in the HLO for my_fft, because x is a 2D array\n  assert(re.search(_PATTERN, pjit_my_fft.lower(x).compile().runtime_executable().hlo_modules()[0].to_string()) is None)\n  # dynamic-slice or all-gather are present in the HLO for fft\n  assert(re.search(_PATTERN, pjit_fft.lower(x).compile().runtime_executable().hlo_modules()[0].to_string())    is not None)\n```\n\nExample:\n```text\n# my_fft\n[[-38.840824   +0.j        -40.649452  +11.845365j\n...\n  -1.6937828  +0.8402481j  15.999859   -4.0156755j]]\n\n# jax.numpy.fft.fft\n[[-38.840824   +0.j        -40.649452  +11.845365j\n  ...\n  -1.6937828  +0.8402481j  15.999859   -4.0156755j]]\n```\n\nExample:\n```text\nwith Mesh(np.array(jax.devices()), ('x',)):\n  x = np.asarray(np.random.randn(32*1024*1024), dtype=np.complex64)\n  y = pjit(lambda x: x, in_shardings=None, out_shardings=P('x'))(x)\n  pjit_my_fft = pjit(my_fft, in_shardings=P('x'), out_shardings=P('x'))\n  pjit_fft    = pjit(fft,    in_shardings=P('x'), out_shardings=P('x'))\n  print(pjit_my_fft(y))\n  print(pjit_fft(y))\n  # dynamic-slice or all-gather are present in the HLO for my_fft, because x is a 1D array\n  assert(re.search(_PATTERN, pjit_my_fft.lower(x).compile().runtime_executable().hlo_modules()[0].to_string()) is None)\n  # dynamic-slice or all-gather are present in the HLO for fft\n  assert(re.search(_PATTERN, pjit_fft.lower(x).compile().runtime_executable().hlo_modules()[0].to_string())    is not None)\n```\n\nExample:\n```text\n# my_fft\n[    7.217285   +0.j     -3012.4937  +4287.635j   -405.83594 +3042.984j\n...  1422.4502  +7271.4297j  -405.84033 -3042.983j\n-3012.4963  -4287.6343j]\n\n# jax.numpy.fft.fft\n[    7.217285   +0.j     -3012.4937  +4287.635j   -405.83594 +3042.984j\n...  1422.4502  +7271.4297j  -405.84033 -3042.983j\n-3012.4963  -4287.6343j]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.801Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":6,"totalLines":134,"estimatedTokens":1289}}87{"id":"doc-jax_experimental_key_reuse_module_jax_documentat-0ef19fcb","source":"documentation","title":"jax.experimental.key_reuse module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.experimental.key_reuse.html","text":"Example:\n```text\n>>> jax.config.update('jax_debug_key_reuse', True)\n```\n\nExample:\n```text\n>>> import jax\n>>> with jax.debug_key_reuse(True):\n...   key = jax.random.key(0)\n...   val1 = jax.random.normal(key)\n...   val2 = jax.random.normal(key)  \nTraceback (most recent call last):\n ...\nKeyReuseError: Previously-consumed key passed to jit-compiled function at index 0\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.801Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":2,"totalLines":18,"estimatedTokens":96}}88{"id":"doc-jax_example_libraries_optimizers_module_jax_docu-1b8d6c58","source":"documentation","title":"jax.example_libraries.optimizers module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.example_libraries.optimizers.html","text":"Example:\n```text\ninit_fun(params)\n\nArgs:\n  params: pytree representing the initial parameters.\n\nReturns:\n  A pytree representing the initial optimizer state, which includes the\n  initial parameters and may also include auxiliary values like initial\n  momentum. The optimizer state pytree structure generally differs from that\n  of `params`.\n```\n\nExample:\n```text\nupdate_fun(step, grads, opt_state)\n\nArgs:\n  step: integer representing the step index.\n  grads: a pytree with the same structure as `get_params(opt_state)`\n    representing the gradients to be used in updating the optimizer state.\n  opt_state: a pytree representing the optimizer state to be updated.\n\nReturns:\n  A pytree with the same structure as the `opt_state` argument representing\n  the updated optimizer state.\n```\n\nExample:\n```text\nget_params(opt_state)\n\nArgs:\n  opt_state: pytree representing an optimizer state.\n\nReturns:\n  A pytree representing the parameters extracted from `opt_state`, such that\n  the invariant `params == get_params(init_fun(params))` holds true.\n```\n\nExample:\n```text\nopt_init, opt_update, get_params = optimizers.sgd(learning_rate)\nopt_state = opt_init(params)\n\ndef step(step, opt_state):\n  value, grads = jax.value_and_grad(loss_fn)(get_params(opt_state))\n  opt_state = opt_update(step, grads, opt_state)\n  return value, opt_state\n\nfor i in range(num_steps):\n  value, opt_state = step(i, opt_state)\n```\n\nExample:\n```text\ninit_fun :: ndarray -> OptStatePytree ndarray\nupdate_fun :: OptStatePytree ndarray -> OptStatePytree ndarray\nget_params :: OptStatePytree ndarray -> ndarray\n```\n\nExample:\n```text\ninit_fun :: ParameterPytree ndarray -> OptimizerState\nupdate_fun :: OptimizerState -> OptimizerState\nget_params :: OptimizerState -> ParameterPytree ndarray\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.803Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":6,"totalLines":70,"estimatedTokens":443}}89{"id":"doc-jax_experimental_sparse_bcoo_jax_documentation-18627b6e","source":"documentation","title":"jax.experimental.sparse.BCOO — JAX documentation","url":"https://docs.jax.dev/en/latest/_autosummary/jax.experimental.sparse.BCOO.html","text":"Example:\n```text\n>>> M = jnp.array([[0., 2., 0.], [1., 0., 4.]])\n>>> M_sp = BCOO.fromdense(M)\n>>> M_sp\nBCOO(float32[2, 3], nse=3)\n```\n\nExample:\n```text\n>>> M_sp.data\nArray([2., 1., 4.], dtype=float32)\n>>> M_sp.indices\nArray([[0, 1],\n       [1, 0],\n       [1, 2]], dtype=int32)\n```\n\nExample:\n```text\n>>> M_sp.todense()\nArray([[0., 2., 0.],\n       [1., 0., 4.]], dtype=float32)\n```\n\nExample:\n```text\n>>> data = jnp.array([1., 3., 5.])\n>>> indices = jnp.array([[0, 0],\n...                      [1, 1],\n...                      [2, 2]])\n>>> mat = BCOO((data, indices), shape=(3, 3))\n>>> mat\nBCOO(float32[3, 3], nse=3)\n>>> mat.todense()\nArray([[1., 0., 0.],\n       [0., 3., 0.],\n       [0., 0., 5.]], dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.804Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":41,"estimatedTokens":182}}90{"id":"doc-jax_experimental_sparse_module_jax_documentation-e6428078","source":"documentation","title":"jax.experimental.sparse module — JAX documentation","url":"https://docs.jax.dev/en/latest/jax.experimental.sparse.html","text":"Example:\n```text\n>>> from jax.experimental import sparse\n>>> import jax.numpy as jnp\n>>> import numpy as np\n```\n\nExample:\n```text\n>>> M = jnp.array([[0., 1., 0., 2.],\n...                [3., 0., 0., 0.],\n...                [0., 0., 4., 0.]])\n```\n\nExample:\n```text\n>>> M_sp = sparse.BCOO.fromdense(M)\n```\n\nExample:\n```text\n>>> M_sp\nBCOO(float32[3, 4], nse=4)\n```\n\nExample:\n```text\n>>> M_sp.todense()\nArray([[0., 1., 0., 2.],\n       [3., 0., 0., 0.],\n       [0., 0., 4., 0.]], dtype=float32)\n```\n\nExample:\n```text\n>>> M_sp.data  # Explicitly stored data\nArray([1., 2., 3., 4.], dtype=float32)\n```\n\nExample:\n```text\n>>> M_sp.indices # Indices of the stored data\nArray([[0, 1],\n       [0, 3],\n       [1, 0],\n       [2, 2]], dtype=int32)\n```\n\nExample:\n```text\n>>> M_sp.ndim\n2\n```\n\nExample:\n```text\n>>> M_sp.shape\n(3, 4)\n```\n\nExample:\n```text\n>>> M_sp.dtype\ndtype('float32')\n```\n\nExample:\n```text\n>>> M_sp.nse  # \"number of specified elements\"\n4\n```\n\nExample:\n```text\n>>> y = jnp.array([3., 6., 5.])\n```\n\nExample:\n```text\n>>> M_sp.T @ y\nArray([18.,  3., 20.,  6.], dtype=float32)\n```\n\nExample:\n```text\n>>> M.T @ y  # Compare to dense version\nArray([18.,  3., 20.,  6.], dtype=float32)\n```\n\nExample:\n```text\n>>> from jax import grad, jit\n```\n\nExample:\n```text\n>>> def f(y):\n...   return (M_sp.T @ y).sum()\n...\n>>> jit(grad(f))(y)\nArray([3., 3., 4.], dtype=float32)\n```\n\nExample:\n```text\n>>> def f(M, v):\n...   return 2 * jnp.dot(jnp.log1p(M.T), v) + 1\n...\n>>> f(M, y)\nArray([17.635532,  5.158883, 17.09438 ,  7.591674], dtype=float32)\n```\n\nExample:\n```text\n>>> f_sp = sparse.sparsify(f)\n```\n\nExample:\n```text\n>>> f_sp(M_sp, y)\nArray([17.635532,  5.158883, 17.09438 ,  7.591674], dtype=float32)\n```\n\nExample:\n```text\n>>> import functools\n>>> from sklearn.datasets import make_classification\n>>> from jax.scipy import optimize\n```\n\nExample:\n```text\n>>> def sigmoid(x):\n...   return 0.5 * (jnp.tanh(x / 2) + 1)\n...\n>>> def y_model(params, X):\n...   return sigmoid(jnp.dot(X, params[1:]) + params[0])\n...\n>>> def loss(params, X, y):\n...   y_hat = y_model(params, X)\n...   return -jnp.mean(y * jnp.log(y_hat) + (1 - y) * jnp.log(1 - y_hat))\n...\n>>> def fit_logreg(X, y):\n...   params = jnp.zeros(X.shape[1] + 1)\n...   result = optimize.minimize(functools.partial(loss, X=X, y=y),\n...                              x0=params, method='BFGS')\n...   return result.x\n```\n\nExample:\n```text\n>>> X, y = make_classification(n_classes=2, random_state=1701)\n>>> params_dense = fit_logreg(X, y)\n>>> print(params_dense)  \n[-0.7298445   0.29893667  1.0248291  -0.44436368  0.8785025  -0.7724008\n -0.62893456  0.2934014   0.82974285  0.16838408 -0.39774987 -0.5071844\n  0.2028872   0.5227761  -0.3739224  -0.7104083   2.4212713   0.6310087\n -0.67060554  0.03139788 -0.05359547]\n```\n\nExample:\n```text\n>>> Xsp = sparse.BCOO.fromdense(X)  # Sparse version of the input\n>>> fit_logreg_sp = sparse.sparsify(fit_logreg)  # Sparse-transformed fit function\n>>> params_sparse = fit_logreg_sp(Xsp, y)\n>>> print(params_sparse)  \n[-0.72971725  0.29878938  1.0246326  -0.44430563  0.8784217  -0.77225566\n -0.6288222   0.29335397  0.8293481   0.16820715 -0.39764675 -0.5069753\n  0.202579    0.522672   -0.3740134  -0.7102678   2.4209507   0.6310593\n -0.670236    0.03132951 -0.05356663]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.807Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":23,"totalLines":173,"estimatedTokens":816}}91{"id":"doc-buffer_donation_jax_documentation-489c5b7c","source":"documentation","title":"Buffer donation — JAX documentation","url":"https://docs.jax.dev/en/latest/buffer_donation.html","text":"Example:\n```text\nparams, state = jax.pmap(update_fn, donate_argnums=(0, 1))(params, state)\n```\n\nExample:\n```text\ndef add(x, y):\n  return x + y\n\nx = jax.device_put(np.ones((2, 3)))\ny = jax.device_put(np.ones((2, 3)))\n# Execute `add` with donation of the buffer for `y`. The result has\n# the same shape and type as `y`, so it will share its buffer.\nz = jax.jit(add, donate_argnums=(1,))(x, y)\n```\n\nExample:\n```text\nparams, state = jax.pmap(update_fn, donate_argnums=(0, 1))(params=params, state=state)\n```\n\nExample:\n```text\ndef add_ones(xs: List[Array]):\n  return [x + 1 for x in xs]\n\nxs = [jax.device_put(np.ones((2, 3))), jax.device_put(np.ones((3, 4)))]\n# Execute `add_ones` with donation of all the buffers for `xs`.\n# The outputs have the same shape and type as the elements of `xs`,\n# so they will share those buffers.\nz = jax.jit(add_ones, donate_argnums=0)(xs)\n```\n\nExample:\n```text\n# Donate the buffer for `y`\nz = jax.jit(add, donate_argnums=(1,))(x, y)\nw = y + 1  # Reuses `y` whose buffer was donated above\n# >> RuntimeError: Invalid argument: CopyToHostAsync() called on invalid buffer\n```\n\nExample:\n```text\n# Execute `add` with donation of the buffers for both `x` and `y`.\n# One of those buffers will be used for the result, but the other will\n# not be used.\nz = jax.jit(add, donate_argnums=(0, 1))(x, y)\n# >> UserWarning: Some donated buffers were not usable: f32[2,3]{1,0}\n```\n\nExample:\n```text\ny = jax.device_put(np.ones((1, 3)))  # `y` has different shape than the output\n# Execute `add` with donation of the buffer for `y`.\nz = jax.jit(add, donate_argnums=(1,))(x, y)\n# >> UserWarning: Some donated buffers were not usable: f32[1,3]{1,0}\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.807Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":7,"totalLines":60,"estimatedTokens":418}}92{"id":"doc-debugging_slow_jax_tracing_and_xla_compilation_j-236c2ff8","source":"documentation","title":"Debugging slow JAX tracing and XLA compilation — JAX documentation","url":"https://docs.jax.dev/en/latest/debugging/slow_tracing_compilation.html","text":"Example:\n```text\nimport jax\njax.config.update(\"jax_log_compiles\", True)\njax.config.update(\"jax_explain_cache_misses\", True)\njax.config.update(\"jax_dump_ir_to\", \"/tmp/jax_ir\")  # or \"sponge\"\njax.config.update(\"jax_dump_ir_modes\", \"eqn_count_pprof\")\n```\n\nExample:\n```text\n# Or via environment variables\nJAX_LOG_COMPILES=1 JAX_EXPLAIN_CACHE_MISSES=1 JAX_DUMP_IR_TO=/tmp/jax_ir JAX_DUMP_IR_MODES=eqn_count_pprof python my_script.py\n```\n\nExample:\n```text\nexport TF_CPP_MIN_LOG_LEVEL=0\n```\n\nExample:\n```text\nW0610 23:50:26.227175 dispatch.py:205] Finished tracing flax_forward for jit in 0.057370424 sec\nW0610 23:50:26.240696 pxla.py:703] Compiling jit(flax_forward) with global shapes and types (ShapedArray(bfloat16[16,16,2048]),). Argument mapping: (UnspecifiedValue,).\nW0610 23:50:26.521261 dispatch.py:205] Finished jaxpr to MLIR module conversion jit(flax_forward) in 0.280440092 sec\nW0610 23:50:44.162414 pxla.py:1234] Finished XLA compilation of jit(flax_forward) in 10.1791505315 sec\n```\n\nExample:\n```text\nI0610 23:32:22.510689 isa_program_util_common.cc:346] (HLO module jit_convert_element_type): Executable fingerprint...\nI0610 23:32:22.561443 isa_program_util_common.cc:346] (HLO module jit_iota): Executable fingerprint...\nI0610 23:32:22.665594 isa_program_util_common.cc:346] (HLO module jit_add): Executable fingerprint...\n```\n\nExample:\n```text\nW0610 23:49:11.800984 partial_eval.py:2179] TRACING CACHE MISS at my_script.py:65:8 (MyLayer.__call__):\n  never seen function:\n    einsum id=55399261860256 defined at aqt_dot_general.py:229\n```\n\nExample:\n```text\nI0610 23:33:55.090166 deepsea_compiler_hlo_passes.cc:6157] HLO_PASSES stage duration: 5.4901s\nI0610 23:33:57.655107 deepsea_compiler_base.cc:3735] BACKEND_PASSES stage duration: 2.5636s\nI0610 23:33:58.089845 deepsea_compiler_backend.cc:1491] CODE_GENERATION stage duration: 434.59ms\nI0610 23:33:59.678832 deepsea_compiler_base.cc:984] END_TO_END stage duration: 10.0876s\n```\n\nExample:\n```text\n# ❌ BAD: Recreating function object on every call\ndef top_function(...):\n  # Creates a brand new function object with a new memory id on every call\n  # to `top_function`.\n  def custom_einsum(a, b):\n    return jnp.einsum(\"...i,...i->...\", a, b)\n\n  return jax.jit(custom_einsum)(x, y)\n```\n\nExample:\n```text\nW0610 23:49:19.396763 partial_eval.py:2179] TRACING CACHE MISS at aqt_flax.py:651:8 (top_function):\n  never seen function:\n    custom_einsum id=55399361734560 defined at aqt_dot_general.py:229\n  but seen another function defined on the same line; maybe the function is\n  being re-defined repeatedly, preventing caching?\n```\n\nExample:\n```text\n# ✔️ GOOD: Reusing the same function handle\ndef custom_einsum(a, b):\n  return jnp.einsum(\"...i,...i->...\", a, b)\n\ndef top_function(...):\n  return jax.jit(custom_einsum)(x, y)\n```\n\nExample:\n```text\n# ❌ BAD: JIT compiling a freshly created lambda on every call\ndef add_multiply(a, b, scale):\n  return (a + b) * scale\n\ndef top_function(x, y, scale_factor):\n  # Creates a brand new lambda object with a new memory id on every pass!\n  return jax.jit(lambda a, b: add_multiply(a, b, scale=scale_factor))(x, y)\n```\n\nExample:\n```text\n# ✔️ GOOD: Using functools.partial\nimport functools\n\ndef add_multiply(a, b, scale):\n  return (a + b) * scale\n\ndef top_function(x, y, scale_factor):\n  # JAX correctly unwraps functools.partial and hits the tracing cache!\n  return jax.jit(functools.partial(add_multiply, scale=scale_factor))(x, y)\n```\n\nExample:\n```text\n# ❌ BAD: Eager op-by-op PyTree mapping\ndef quantize(x):\n  scale = jnp.max(jnp.abs(x))\n  return jnp.round(x * scale)\n\n# Eagerly dispatches jnp.max, jnp.abs, jnp.round for every weight tensor!\nparams = jax.tree.map(quantize, params)\n```\n\nExample:\n```text\n# ✔️ GOOD: Compiles a single fused XLA graph for the entire PyTree\nparams = jax.jit(lambda p: jax.tree.map(quantize, p))(params)\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\n\n@jax.jit\ndef unrolled_loop(x):\n  # ❌ BAD: Unrolls 5,000 identical add equations into the Jaxpr!\n  for _ in range(5000):\n    x = x + 1.0\n  return x\n\n# Run with JAX_DUMP_IR_TO=/tmp/jax_ir JAX_DUMP_IR_MODES=eqn_count_pprof\nunrolled_loop(jnp.zeros(10))\n```\n\nExample:\n```text\nshowing top 5 nodes out of 5\n      flat  flat%   sum%        cum   cum%\n      5000 99.98% 99.98%       5000 99.98%  my_script.py:8 (unrolled_loop)\n         1  0.02% 100.0%          1  0.02%  my_script.py:12 (<module>)\n```\n\nExample:\n```text\npprof -http=localhost:8080 /tmp/jax_ir/jax_000001_unrolled_loop.eqn_count_pprof\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.808Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":17,"totalLines":152,"estimatedTokens":1125}}93{"id":"doc-transfer_guard_jax_documentation-93c3e275","source":"documentation","title":"Transfer guard — JAX documentation","url":"https://docs.jax.dev/en/latest/transfer_guard.html","text":"Example:\n```text\n>>> jax.config.update(\"jax_transfer_guard\", \"allow\")  # This is default.\n>>>\n>>> x = jnp.array(1)\n>>> y = jnp.array(2)\n>>> z = jnp.array(3)\n>>>\n>>> print(\"x\", x)  # All transfers are allowed.\nx 1\n>>> with jax.transfer_guard(\"disallow\"):\n...   print(\"x\", x)  # x has already been fetched into the host.\n...   print(\"y\", jax.device_get(y))  # Explicit transfers are allowed.\n...   try:\n...     print(\"z\", z)  # Implicit transfers are disallowed.\n...     assert False, \"This line is expected to be unreachable.\"\n...   except:\n...     print(\"z could not be fetched\")  \nx 1\ny 2\nz could not be fetched\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.809Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":1,"totalLines":24,"estimatedTokens":158}}94{"id":"doc-debugging_runtime_values_jax_documentation-8a7ae8c0","source":"documentation","title":"Debugging runtime values — JAX documentation","url":"https://docs.jax.dev/en/latest/debugging/index.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n\n@jax.jit\ndef f(x):\n  jax.debug.print(\"🤯 {x} 🤯\", x=x)\n  y = jnp.sin(x)\n  jax.debug.breakpoint()\n  jax.debug.print(\"🤯 {y} 🤯\", y=y)\n  return y\n\nf(2.)\n# Prints:\n# 🤯 2.0 🤯\n# Enters breakpoint to inspect values!\n# 🤯 0.9092974662780762 🤯\n```\n\nExample:\n```text\nfrom jax.experimental import checkify\nimport jax\nimport jax.numpy as jnp\n\ndef f(x, i):\n  checkify.check(i >= 0, \"index needs to be non-negative!\")\n  y = x[i]\n  z = jnp.sin(y)\n  return z\n\njittable_f = checkify.checkify(f)\n\nerr, z = jax.jit(jittable_f)(jnp.ones((5,)), -1)\nprint(err.get())\n# >> index needs to be non-negative! (check failed at <...>:6 (f))\n```\n\nExample:\n```text\nerrors = checkify.user_checks | checkify.index_checks | checkify.float_checks\nchecked_f = checkify.checkify(f, errors=errors)\n\nerr, z = checked_f(jnp.ones((5,)), 100)\nerr.throw()\n# ValueError: out-of-bounds indexing at <..>:7 (f)\n\nerr, z = checked_f(jnp.ones((5,)), -1)\nerr.throw()\n# ValueError: index needs to be non-negative! (check failed at <…>:6 (f))\n\nerr, z = checked_f(jnp.array([jnp.inf, 1]), 0)\nerr.throw()\n# ValueError: nan generated by primitive sin at <...>:8 (f)\n```\n\nExample:\n```text\nimport jax\njax.config.update(\"jax_debug_nans\", True)\n\ndef f(x, y):\n  return x / y\njax.jit(f)(0., 0.)  # ==> raises FloatingPointError exception!\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.experimental.xla_metadata import set_xla_metadata\n\n# Tagging an individual operation\ndef value_tagging(x):\n  y = jnp.sin(x)\n  z = jnp.cos(x)\n  return set_xla_metadata(y * z, breakpoint=True)\n\nprint(jax.jit(value_tagging).lower(1.0).as_text(\"hlo\"))\n```\n\nExample:\n```text\nENTRY main.5 {\n  x.1 = f32[] parameter(0)\n  sin.2 = f32[] sine(x.1)\n  cos.3 = f32[] cosine(x.1)\n  ROOT mul.4 = f32[] multiply(sin.2, cos.3), frontend_attributes={breakpoint=\"true\"}\n}\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.809Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":6,"totalLines":93,"estimatedTokens":468}}95{"id":"doc-attaching_xla_metadata_with_set_xla_metadata_jax-f8400b9c","source":"documentation","title":"Attaching XLA Metadata with set_xla_metadata — JAX documentation","url":"https://docs.jax.dev/en/latest/debugging/xla_metadata.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.experimental.xla_metadata import set_xla_metadata\n\n# Tagging an individual operation\ndef value_tagging(x):\n  y = jnp.sin(x)\n  z = jnp.cos(x)\n  return set_xla_metadata(y * z, breakpoint=True)\n\nprint(jax.jit(value_tagging).lower(1.0).as_text(\"hlo\"))\n```\n\nExample:\n```text\nENTRY main.5 {\n  x.1 = f32[] parameter(0)\n  sin.2 = f32[] sine(x.1)\n  cos.3 = f32[] cosine(x.1)\n  ROOT mul.4 = f32[] multiply(sin.2, cos.3), frontend_attributes={breakpoint=\"true\"}\n}\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.experimental.xla_metadata import set_xla_metadata\n\n# Tagging a block of code\ndef context_tagging(x):\n  with set_xla_metadata(_xla_log=True):\n    y = jnp.sin(x)\n    z = jnp.cos(y)\n    return y * z\n\nprint(jax.jit(context_tagging).lower(1.0).as_text(\"hlo\"))\n```\n\nExample:\n```text\nENTRY main.5 {\n  x.1 = f32[] parameter(0)\n  sin.2 = f32[] sine(x.1), frontend_attributes={_xla_log=\"true\"}\n  cos.3 = f32[] cosine(sin.2), frontend_attributes={_xla_log=\"true\"}\n  ROOT mul.4 = f32[] multiply(sin.2, cos.3), frontend_attributes={_xla_log=\"true\"}\n}\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.experimental.xla_metadata import set_xla_metadata\n\n# Tagging with a decorator\n@set_xla_metadata(_xla_log=True)\n@jax.jit\ndef decorator_tagging(x):\n  y = jnp.sin(x)\n  z = jnp.cos(y)\n  return y * z\n\nprint(decorator_tagging.lower(1.0).as_text(\"hlo\"))\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.experimental.xla_metadata import set_xla_metadata\n\ndef fn(x):\n    y = jnp.sin(x)\n    z = jnp.cos(x)\n    return y * z\n\nmetadata = {\"example\": \"grad_tagging\"}\n\n# --- Define Custom VJP to tag gradients ---\n@jax.custom_vjp\ndef wrapped_fn(x):\n    return fn(x)\n\ndef fwd(*args):\n    primal_out, vjp_fn = jax.vjp(fn, *args)\n    return primal_out, vjp_fn\n\ndef bwd(vjp_fn, cts_in):\n    cts_out = vjp_fn(cts_in)\n    cts_out = set_xla_metadata(cts_out, **metadata)\n    return cts_out\n\nwrapped_fn.defvjp(fwd, bwd)\n# ------\n\nprint(jax.jit(jax.grad(wrapped_fn)).lower(jnp.array(3.0)).as_text(\"hlo\"))\n```\n\nExample:\n```text\nENTRY main.10 {\n  x.1 = f32[] parameter(0)\n  sin.2 = f32[] sine(x.1)\n  neg.6 = f32[] negate(sin.2)\n  sin.5 = f32[] sine(x.1)\n  mul.7 = f32[] multiply(neg.6, sin.5)\n  cos.4 = f32[] cosine(x.1)\n  cos.3 = f32[] cosine(x.1)\n  mul.8 = f32[] multiply(cos.4, cos.3)\n  ROOT add_any.9 = f32[] add(mul.7, mul.8), frontend_attributes={example=\"grad_tagging\"}\n}\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.810Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":7,"totalLines":117,"estimatedTokens":618}}96{"id":"doc-profiling_device_memory_jax_documentation-451769d1","source":"documentation","title":"Profiling device memory — JAX documentation","url":"https://docs.jax.dev/en/latest/device_memory_profiling.html","text":"Example:\n```text\ngo install github.com/google/pprof@latest\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nimport jax.profiler\n\ndef func1(x):\n  return jnp.tile(x, 10) * 0.5\n\ndef func2(x):\n  y = func1(x)\n  return y, jnp.tile(x, 10) + 1\n\nx = jax.random.normal(jax.random.key(42), (1000, 1000))\ny, z = func2(x)\n\nz.block_until_ready()\n\njax.profiler.save_device_memory_profile(\"memory.prof\")\n```\n\nExample:\n```text\npprof --web memory.prof\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nimport jax.profiler\n\ndef afunction():\n  return jax.random.normal(jax.random.key(77), (1000000,))\n\nz = afunction()\n\ndef anotherfunc():\n  arrays = []\n  for i in range(1, 10):\n    x = jax.random.normal(jax.random.key(42), (i, 10000))\n    arrays.append(x)\n    x.block_until_ready()\n    jax.profiler.save_device_memory_profile(f\"memory{i}.prof\")\n\nanotherfunc()\n```\n\nExample:\n```text\npprof --web memory9.prof\n```\n\nExample:\n```text\npprof --web --diff_base memory1.prof memory9.prof\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.810Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":6,"totalLines":64,"estimatedTokens":248}}97{"id":"doc-jax_debugging_flags_jax_documentation-137d7823","source":"documentation","title":"JAX debugging flags — JAX documentation","url":"https://docs.jax.dev/en/latest/debugging/flags.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\nimport traceback\njax.config.update(\"jax_debug_nans\", True)\n\ndef f(x):\n  w = 3 * jnp.square(x)\n  return jnp.log(-w)\n\n# The stack trace is very long so only print a couple lines.\ntry:\n  f(5.)\nexcept FloatingPointError as e:\n  print(traceback.format_exc(limit=2))\n```\n\nExample:\n```text\nInvalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.\nTraceback (most recent call last):\n  File \"/tmp/ipykernel_1711/1479925735.py\", line 12, in <module>\n    f(5.)\n  File \"/tmp/ipykernel_1711/1479925735.py\", line 8, in f\n    return jnp.log(-w)\n           ^^^^^^^^^^^\nFloatingPointError: invalid value (nan) encountered in log\n```\n\nExample:\n```text\njax.jit(f)(5.)\n```\n\nExample:\n```text\nInvalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.\nInvalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.\n```\n\nExample:\n```text\n---------------------------------------------------------------------------\nFloatingPointError                        Traceback (most recent call last)\nCell In[2], line 1\n----> 1 jax.jit(f)(5.)\n\n    [... skipping hidden 5 frame]\n\nCell In[1], line 8, in f(x)\n      6 def f(x):\n      7   w = 3 * jnp.square(x)\n----> 8   return jnp.log(-w)\n\n    [... skipping hidden 5 frame]\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/ufuncs.py:491, in log(x)\n    456 @export\n    457 @jit(inline=True)\n    458 def log(x: ArrayLike, /) -> Array:\n    459   \"\"\"Calculate element-wise natural logarithm of the input.\n    460 \n    461   JAX implementation of :obj:`numpy.log`.\n   (...)    489     Array(True, dtype=bool)\n    490   \"\"\"\n--> 491   out = lax.log(*promote_args_inexact('log', x))\n    492   jnp_error._set_error_if_nan(out)\n    493   return out\n\n    [... skipping hidden 7 frame]\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/pjit.py:175, in _run_python_pjit(p, args_flat, fun, args, kwargs)\n    173 except api_util.InternalFloatingPointError as e:\n    174   if getattr(fun, '_apply_primitive', False):\n--> 175     raise FloatingPointError(\n    176         f\"invalid value ({e.ty}) encountered in {fun.__qualname__}\") from None\n    177   api_util.maybe_recursive_nan_check(e, fun, args, kwargs)  # should always raise.\n    178   raise RuntimeError(\"Internal error\") from e  # fall-back error to be safe.\n\nFloatingPointError: invalid value (nan) encountered in log\n```\n\nExample:\n```text\nwith jax.debug_nans(False):\n  print(jax.jit(f)(5.))\n```\n\nExample:\n```text\nnan\n```\n\nExample:\n```text\nimport jax\njax.config.update(\"jax_disable_jit\", True)\n\ndef f(x):\n  y = jnp.log(x)\n  if jnp.isnan(y):\n    breakpoint()\n  return y\njax.jit(f)(-2.)  # ==> Enters PDB breakpoint!\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.811Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":8,"totalLines":108,"estimatedTokens":718}}98{"id":"doc-profiling_computation_jax_documentation-91b3d305","source":"documentation","title":"Profiling computation — JAX documentation","url":"https://docs.jax.dev/en/latest/profiling.html","text":"Example:\n```text\nwith jax.profiler.trace(\"/tmp/jax-trace\", create_perfetto_link=True):\n  # Run the operations to be profiled\n  key = jax.random.key(0)\n  x = jax.random.normal(key, (5000, 5000))\n  y = x @ x\n  y.block_until_ready()\n```\n\nExample:\n```text\n$ ssh -L 9001:127.0.0.1:9001 <user>@<host>\n```\n\nExample:\n```text\n$ gcloud compute ssh <machine-name> -- -L 9001:127.0.0.1:9001\n```\n\nExample:\n```text\n$ python -m jax.collect_profile <port> <duration_in_ms>\n```\n\nExample:\n```text\npip install xprof\n```\n\nExample:\n```text\npip install tb-nightly xprof-nightly\n```\n\nExample:\n```text\n$ tensorboard --logdir=/tmp/profile-data\n[...]\nServing TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all\nTensorBoard 2.19.0 at http://localhost:6006/ (Press CTRL+C to quit)\n```\n\nExample:\n```text\nimport jax\n\njax.profiler.start_trace(\"/tmp/profile-data\")\n\n# Run the operations to be profiled\nkey = jax.random.key(0)\nx = jax.random.normal(key, (5000, 5000))\ny = x @ x\ny.block_until_ready()\n\njax.profiler.stop_trace()\n```\n\nExample:\n```text\nimport jax\n\nwith jax.profiler.trace(\"/tmp/profile-data\"):\n  key = jax.random.key(0)\n  x = jax.random.normal(key, (5000, 5000))\n  y = x @ x\n  y.block_until_ready()\n```\n\nExample:\n```text\n$ xprof --port 8791 /tmp/profile-data\nAttempting to start XProf server:\n  Log Directory: /tmp/profile-data\n  Port: 8791\nXProf at http://localhost:8791/ (Press CTRL+C to quit)\n```\n\nExample:\n```text\nxprof --logdir /tmp/profile-data/\n```\n\nExample:\n```text\nimport jax.profiler\njax.profiler.start_server(9999)\n```\n\nExample:\n```text\nimport jax\n\noptions = jax.profiler.ProfileOptions()\noptions.python_tracer_level = 0\noptions.host_tracer_level = 0\njax.profiler.start_trace(\"/tmp/profile-data\", profiler_options=options)\n\n# Run the operations to be profiled\nkey = jax.random.key(0)\nx = jax.random.normal(key, (5000, 5000))\ny = x @ x\ny.block_until_ready()\n\njax.profiler.stop_trace()\n```\n\nExample:\n```text\noptions = ProfileOptions()\noptions.advanced_configuration = {\"tpu_trace_mode\" : \"TRACE_ONLY_HOST\", \"tpu_num_sparse_cores_to_trace\" : 2}\n```\n\nExample:\n```text\nW external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libcupti.so.10.1'; dlerror: libcupti.so.10.1: cannot open shared object file: No such file or directory\n2020-06-12 13:19:59.822799: E external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1422] function cupti_interface_->Subscribe( &subscriber_, (CUpti_CallbackFunc)ApiCallback, this)failed with error CUPTI could not be loaded or symbol could not be found.\n```\n\nExample:\n```text\nexport LD_LIBRARY_PATH=/usr/local/cuda-10.1/extras/CUPTI/lib64/:$LD_LIBRARY_PATH\n```\n\nExample:\n```text\nE external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1445] function cupti_interface_->EnableCallback( 0 , subscriber_, CUPTI_CB_DOMAIN_DRIVER_API, cbid)failed with error CUPTI_ERROR_INSUFFICIENT_PRIVILEGES\n2020-06-12 14:31:54.097791: E external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1487] function cupti_interface_->ActivityDisable(activity)failed with error CUPTI_ERROR_NOT_INITIALIZED\n```\n\nExample:\n```text\necho 'options nvidia \"NVreg_RestrictProfilingToAdminUsers=0\"' | sudo tee -a /etc/modprobe.d/nvidia-kernel-common.conf\nsudo update-initramfs -u\nsudo reboot now\n```\n\nExample:\n```text\nssh -L 6006:localhost:6006 <remote server address>\n```\n\nExample:\n```text\n$ gcloud compute ssh <machine-name> -- -L 6006:localhost:6006\n```\n\nExample:\n```text\npip uninstall tensorflow tf-nightly tensorboard tb-nightly xprof xprof-nightly tensorboard-plugin-profile tbp-nightly\npip install tensorboard xprof\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.812Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":21,"totalLines":154,"estimatedTokens":921}}99{"id":"doc-benchmarking_jax_code_jax_documentation-be21e3bd","source":"documentation","title":"Benchmarking JAX code — JAX documentation","url":"https://docs.jax.dev/en/latest/benchmarking.html","text":"Example:\n```text\nimport numpy as np\nimport jax\n\ndef f(x):  # function we're benchmarking (works in both NumPy & JAX)\n    return x.T @ (x - x.mean(axis=0))\n\nx_np = np.ones((1000, 1000), dtype=np.float32)  # same as JAX default dtype\n%timeit f(x_np)  # measure NumPy runtime\n\n# measure JAX device transfer time\n%time x_jax = jax.device_put(x_np).block_until_ready()\n\nf_jit = jax.jit(f)\n%time f_jit(x_jax).block_until_ready()  # measure JAX compilation time\n%timeit f_jit(x_jax).block_until_ready()  # measure JAX runtime\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.812Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":1,"totalLines":20,"estimatedTokens":134}}100{"id":"doc-compiled_prints_and_breakpoints_jax_documentatio-152aa46b","source":"documentation","title":"Compiled prints and breakpoints — JAX documentation","url":"https://docs.jax.dev/en/latest/debugging/print_breakpoint.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n\n@jax.jit\ndef f(x):\n  jax.debug.print(\"🤯 {x} 🤯\", x=x)\n  y = jnp.sin(x)\n  jax.debug.print(\"🤯 {y} 🤯\", y=y)\n  return y\n\nf(2.)\n# Prints:\n# 🤯 2.0 🤯\n# 🤯 0.9092974662780762 🤯\n```\n\nExample:\n```text\ndef debug.print(fmt: str, *args: PyTree[Array], **kwargs: PyTree[Array]) -> None:\n  print(fmt.format(*args, **kwargs))\n```\n\nExample:\n```text\nxs = jnp.arange(3.)\n\ndef f(x):\n  jax.debug.print(\"x: {}\", x)\n  y = jnp.sin(x)\n  jax.debug.print(\"y: {}\", y)\n  return y\njax.vmap(f)(xs)\n# Prints: x: 0.0\n#         x: 1.0\n#         x: 2.0\n#         y: 0.0\n#         y: 0.841471\n#         y: 0.9092974\njax.lax.map(f, xs)\n# Prints: x: 0.0\n#         y: 0.0\n#         x: 1.0\n#         y: 0.841471\n#         x: 2.0\n#         y: 0.9092974\n```\n\nExample:\n```text\nxs = jnp.arange(2.)\n\ndef f(x):\n  jax.debug.print(\"x: {}\", x)\n  return x\njax.pmap(f)(xs)\n# Prints: x: 0.0\n#         x: 1.0\n# OR\n# Prints: x: 1.0\n#         x: 0.0\n```\n\nExample:\n```text\ndef f(x):\n  jax.debug.print(\"x: {}\", x)\n  return x * 2.\n\njax.grad(f)(1.)\n# Prints: x: 1.0\n```\n\nExample:\n```text\n@jax.custom_vjp\ndef print_grad(x):\n  return x\n\ndef print_grad_fwd(x):\n  return x, None\n\ndef print_grad_bwd(_, x_grad):\n  jax.debug.print(\"x_grad: {}\", x_grad)\n  return (x_grad,)\n\nprint_grad.defvjp(print_grad_fwd, print_grad_bwd)\n\n\ndef f(x):\n  x = print_grad(x)\n  return x * 2.\njax.grad(f)(1.)\n# Prints: x_grad: 2.0\n```\n\nExample:\n```text\ndef callback(fun: Callable, *args: PyTree[Array], **kwargs: PyTree[Array]) -> None:\n  fun(*args, **kwargs)\n  return None\n```\n\nExample:\n```text\n@jax.jit\ndef f(x, y):\n  jax.debug.print(\"x: {}\", x)\n  jax.debug.print(\"y: {}\", y)\n  return x + y\n\nf(2., 3.)\n# Prints: x: 2.0\n#         y: 3.0\n# OR\n# Prints: y: 3.0\n#         x: 2.0\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  jax.debug.print(\"x: {}\", x)\n  return x\nf(2.).block_until_ready()\n# <do something else>\n# Prints: x: 2.\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  jax.debug.print(\"x: {}\", x)\n  return x\nf(2.).block_until_ready()\njax.effects_barrier()\n# Prints: x: 2.\n# <do something else>\n```\n\nExample:\n```text\ndef f(w, b, x):\n  logits = w.dot(x) + b\n  jax.debug.print(\"logits: {}\", logits)\n  return jax.nn.relu(logits)\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  y, z = jnp.sin(x), jnp.cos(x)\n  jax.debug.breakpoint()\n  return y * z\nf(2.) # ==> Pauses during execution!\n```\n\nExample:\n```text\ndef breakpoint_if_nonfinite(x):\n  is_finite = jnp.isfinite(x).all()\n  def true_fn(x):\n    pass\n  def false_fn(x):\n    jax.debug.breakpoint()\n  lax.cond(is_finite, true_fn, false_fn, x)\n\n@jax.jit\ndef f(x, y):\n  z = x / y\n  breakpoint_if_nonfinite(z)\n  return z\nf(2., 0.) # ==> Pauses during execution!\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.813Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":13,"totalLines":180,"estimatedTokens":671}}101{"id":"doc-persistent_compilation_cache_jax_documentation-1cbc6b68","source":"documentation","title":"Persistent compilation cache — JAX documentation","url":"https://docs.jax.dev/en/latest/persistent_compilation_cache.html","text":"Example:\n```text\npip install etils\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\n\njax.config.update(\"jax_compilation_cache_dir\", \"/tmp/jax_cache\")\njax.config.update(\"jax_persistent_cache_min_entry_size_bytes\", -1)\njax.config.update(\"jax_persistent_cache_min_compile_time_secs\", 0)\njax.config.update(\"jax_persistent_cache_enable_xla_caches\", \"xla_gpu_per_fusion_autotune_cache_dir\")\n\n@jax.jit\ndef f(x):\n  return x + 1\n\nx = jnp.zeros((2, 2))\nf(x)\n```\n\nExample:\n```text\nexport JAX_COMPILATION_CACHE_DIR=\"/tmp/jax_cache\"\n```\n\nExample:\n```text\nimport os\nos.environ[\"JAX_COMPILATION_CACHE_DIR\"] = \"/tmp/jax_cache\"\n```\n\nExample:\n```text\njax.config.update(\"jax_compilation_cache_dir\", \"/tmp/jax_cache\")\n```\n\nExample:\n```text\nfrom jax.experimental.compilation_cache import compilation_cache as cc\ncc.set_cache_dir(\"/tmp/jax_cache\")\n```\n\nExample:\n```text\n# Example assuming the GCS bucket is mounted at /gcs/my-bucket\njax.config.update(\"jax_compilation_cache_dir\", \"/gcs/my-bucket/jax-cache\")\n```\n\nExample:\n```text\njax.config.update(\"jax_compilation_cache_dir\", \"gs://jax-cache\")\n```\n\nExample:\n```text\njax.config.update(\"jax_mock_gpu_topology\", \"4x8x1\")\n```\n\nExample:\n```text\nimport os\nos.environ[\"JAX_DEBUG_LOG_MODULES\"] = \"jax._src.compiler,jax._src.lru_cache\"\n```\n\nExample:\n```text\nimport os\nos.environ[\"JAX_LOGGING_LEVEL\"] = \"DEBUG\"\n# or locally with\njax.config.update(\"jax_logging_level\", \"DEBUG\")\n```\n\nExample:\n```text\njax.config.update(\"jax_explain_cache_misses\", True)\n```\n\nExample:\n```text\nimport jax\n\ndef F(x1, x2, gamma, beta):\n   ln_out = LayerNorm(x1, gamma, beta)\n   return ln_out @ x2\n```\n\nExample:\n```text\nlayernorm_matmul_without_shard_map = jax.jit(F, in_shardings=(...), out_sharding=(...))(x1, x2, gamma, beta)\n```\n\nExample:\n```text\nimport jax\n\ndef G(x1, x2, gamma, beta, mesh, ispecs, ospecs):\n   ln_out = jax.shard_map(LayerNorm, mesh=mesh, in_specs=ispecs, out_specs=ospecs, check_vma=False)(x1, x2, gamma, beta)\n   return ln_out @ x2\n\nispecs = jax.sharding.PartitionSpec(...)\nospecs = jax.sharding.PartitionSpec(...)\nmesh = jax.sharding.Mesh(...)\nlayernorm_matmul_with_shard_map = jax.jit(G, static_argnames=['mesh', 'ispecs', 'ospecs'])(x1, x2, gamma, beta, mesh, ispecs, ospecs)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.814Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":15,"totalLines":109,"estimatedTokens":556}}102{"id":"doc-custom_pytree_nodes_jax_documentation-d208dd79","source":"documentation","title":"Custom pytree nodes — JAX documentation","url":"https://docs.jax.dev/en/latest/custom_pytrees.html","text":"Example:\n```text\nimport jax\n\nclass Special(object):\n  def __init__(self, x, y):\n    self.x = x\n    self.y = y\n\njax.tree.leaves([\n    Special(0, 1),\n    Special(2, 4),\n])\n```\n\nExample:\n```text\n[<__main__.Special at 0x708686f03a40>, <__main__.Special at 0x708686f033e0>]\n```\n\nExample:\n```text\njax.tree.map(lambda x: x + 1,\n  [\n    Special(0, 1),\n    Special(2, 4)\n  ])\n```\n\nExample:\n```text\n---------------------------------------------------------------------------\nTypeError                                 Traceback (most recent call last)\nCell In[2], line 1\n----> 1 jax.tree.map(lambda x: x + 1,\n      2   [\n      3     Special(0, 1),\n      4     Special(2, 4)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/tree.py:156, in map(f, tree, is_leaf, *rest)\n    116 def map(f: Callable[..., Any],\n    117         tree: Any,\n    118         *rest: Any,\n    119         is_leaf: Callable[[Any], bool] | None = None) -> Any:\n    120   \"\"\"Maps a multi-input function over pytree args to produce a new pytree.\n    121 \n    122   Args:\n   (...)    154     - :func:`jax.tree.reduce`\n    155   \"\"\"\n--> 156   return tree_util.tree_map(f, tree, *rest, is_leaf=is_leaf)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/tree_util.py:400, in tree_map(f, tree, is_leaf, *rest)\n    398   err = next(_prefix_error((), tree, r2, is_leaf), None)  # type: ignore\n    399   raise (err('tree_map tree') if err is not None else e) from None\n--> 400 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/tree_util.py:400, in <genexpr>(.0)\n    398   err = next(_prefix_error((), tree, r2, is_leaf), None)  # type: ignore\n    399   raise (err('tree_map tree') if err is not None else e) from None\n--> 400 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))\n\nCell In[2], line 1, in <lambda>(x)\n----> 1 jax.tree.map(lambda x: x + 1,\n      2   [\n      3     Special(0, 1),\n      4     Special(2, 4)\n\nTypeError: unsupported operand type(s) for +: 'Special' and 'int'\n```\n\nExample:\n```text\nfrom jax.tree_util import register_pytree_node\n\nclass RegisteredSpecial(Special):\n  def __repr__(self):\n    return \"RegisteredSpecial(x={}, y={})\".format(self.x, self.y)\n\ndef special_flatten(v):\n  \"\"\"Specifies a flattening recipe.\n\n  Params:\n    v: The value of the registered type to flatten.\n  Returns:\n    A pair of an iterable with the children to be flattened recursively,\n    and some opaque auxiliary data to pass back to the unflattening recipe.\n    The auxiliary data is stored in the treedef for use during unflattening.\n    The auxiliary data could be used, for example, for dictionary keys.\n  \"\"\"\n  children = (v.x, v.y)\n  aux_data = None\n  return (children, aux_data)\n\ndef special_unflatten(aux_data, children):\n  \"\"\"Specifies an unflattening recipe.\n\n  Params:\n    aux_data: The opaque data that was specified during flattening of the\n      current tree definition.\n    children: The unflattened children\n\n  Returns:\n    A reconstructed object of the registered type, using the specified\n    children and auxiliary data.\n  \"\"\"\n  return RegisteredSpecial(*children)\n\n# Global registration\nregister_pytree_node(\n    RegisteredSpecial,\n    special_flatten,    # Instruct JAX what are the children nodes.\n    special_unflatten   # Instruct JAX how to pack back into a `RegisteredSpecial`.\n)\n```\n\nExample:\n```text\njax.tree.map(lambda x: x + 1,\n  [\n   RegisteredSpecial(0, 1),\n   RegisteredSpecial(2, 4),\n  ])\n```\n\nExample:\n```text\n[RegisteredSpecial(x=1, y=2), RegisteredSpecial(x=3, y=5)]\n```\n\nExample:\n```text\nfrom jax.tree_util import register_pytree_node_class\n\n@register_pytree_node_class\nclass RegisteredSpecial2(Special):\n  def __repr__(self):\n    return \"RegisteredSpecial2(x={}, y={})\".format(self.x, self.y)\n\n  def tree_flatten(self):\n    children = (self.x, self.y)\n    aux_data = None\n    return (children, aux_data)\n\n  @classmethod\n  def tree_unflatten(cls, aux_data, children):\n    return cls(*children)\n\n\ndef show_example(structured):\n  flat, tree = structured.tree_flatten()\n  unflattened = RegisteredSpecial2.tree_unflatten(tree, flat)\n  print(f\"{structured=}\\n  {flat=}\\n  {tree=}\\n  {unflattened=}\")\n\n\nshow_example(RegisteredSpecial2(1., 2.))\n```\n\nExample:\n```text\nstructured=RegisteredSpecial2(x=1.0, y=2.0)\n  flat=(1.0, 2.0)\n  tree=None\n  unflattened=RegisteredSpecial2(x=1.0, y=2.0)\n```\n\nExample:\n```text\nfrom typing import NamedTuple, Any\n\nclass MyOtherContainer(NamedTuple):\n  name: str\n  a: Any\n  b: Any\n  c: Any\n\n# NamedTuple subclasses are handled as pytree nodes, so\n# this will work out-of-the-box.\njax.tree.leaves([\n    MyOtherContainer('Alice', 1, 2, 3),\n    MyOtherContainer('Bob', 4, 5, 6)\n])\n```\n\nExample:\n```text\n['Alice', 1, 2, 3, 'Bob', 4, 5, 6]\n```\n\nExample:\n```text\nfrom dataclasses import dataclass\nimport jax.numpy as jnp\nimport numpy as np\nimport functools\n\n@functools.partial(jax.tree_util.register_dataclass,\n                   data_fields=['a', 'b', 'c'],\n                   meta_fields=['name'])\n@dataclass\nclass MyDataclassContainer(object):\n  name: str\n  a: Any\n  b: Any\n  c: Any\n\n# MyDataclassContainer is now a pytree node.\njax.tree.leaves([\n  MyDataclassContainer('apple', 5.3, 1.2, jnp.zeros([4])),\n  MyDataclassContainer('banana', np.array([3, 4]), -1., 0.)\n])\n```\n\nExample:\n```text\n[5.3, 1.2, Array([0., 0., 0., 0.], dtype=float32), array([3, 4]), -1.0, 0.0]\n```\n\nExample:\n```text\n@jax.jit\ndef f(x: MyDataclassContainer | MyOtherContainer):\n  return x.a + x.b\n\n# Works fine! `mdc.name` is static.\nmdc = MyDataclassContainer('mdc', 1, 2, 3)\ny = f(mdc)\n```\n\nExample:\n```text\nmoc = MyOtherContainer('moc', 1, 2, 3)\ny = f(moc)\n```\n\nExample:\n```text\n---------------------------------------------------------------------------\nTypeError                                 Traceback (most recent call last)\n    [... skipping hidden 1 frame]\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:1960, in shaped_abstractify(x)\n   1959 if (aval_fn := pytype_aval_mappings.get(typ)):  # fast path\n-> 1960   return aval_fn(x)\n   1961 for t in typ.__mro__[1:]:\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/core.py:948, in _str_abstractify(x)\n    947 def _str_abstractify(x):\n--> 948   raise TypeError(f\"Argument '{x}' of type {type(x)} is not a valid JAX type\")\n\nTypeError: Argument 'moc' of type <class 'str'> is not a valid JAX type\n\nThe above exception was the direct cause of the following exception:\n\nTypeError                                 Traceback (most recent call last)\nCell In[9], line 2\n      1 moc = MyOtherContainer('moc', 1, 2, 3)\n----> 2 y = f(moc)\n\n    [... skipping hidden 3 frame]\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/pjit.py:668, in _infer_input_type(fun, dbg_fn, explicit_args)\n    666   dbg = dbg_fn()\n    667   arg_description = f\"path {dbg.arg_names[i] if dbg.arg_names is not None else 'unknown'}\"\n--> 668   raise TypeError(\n    669     f\"Error interpreting argument to {fun} as an abstract array.\"\n    670     f\" The problematic value is of type {type(x)} and was passed to\"\n    671     f\" the function at {arg_description}.\\n\"\n    672     \"This typically means that a jit-wrapped function was called with a non-array\"\n    673     \" argument, and this argument was not marked as static using the\"\n    674     \" static_argnums or static_argnames parameters of jax.jit.\"\n    675   ) from e\n    676 if config.mutable_array_checks.value:\n    677   check_no_aliased_ref_args(dbg_fn, avals, explicit_args)\n\nTypeError: Error interpreting argument to <function f at 0x70865ea09f80> as an abstract array. The problematic value is of type <class 'str'> and was passed to the function at path x.name.\nThis typically means that a jit-wrapped function was called with a non-array argument, and this argument was not marked as static using the static_argnums or static_argnames parameters of jax.jit.\n```\n\nExample:\n```text\nclass MyTree:\n  def __init__(self, a):\n    self.a = jnp.asarray(a)\n\nregister_pytree_node(MyTree, lambda tree: ((tree.a,), None),\n    lambda _, args: MyTree(*args))\n\ntree = MyTree(jnp.arange(5.0))\n\njax.vmap(lambda x: x)(tree)      # Error because object() is passed to `MyTree`.\n```\n\nExample:\n```text\n<__main__.MyTree at 0x70864f25fad0>\n```\n\nExample:\n```text\njax.jacobian(lambda x: x)(tree)  # Error because MyTree(...) is passed to `MyTree`.\n```\n\nExample:\n```text\n---------------------------------------------------------------------------\nValueError                                Traceback (most recent call last)\nCell In[11], line 1\n----> 1 jax.jacobian(lambda x: x)(tree)  # Error because MyTree(...) is passed to `MyTree`.\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/api.py:837, in jacrev.<locals>.jacfun(*args, **kwargs)\n    834 @wraps(fun, docstr=docstr, argnums=argnums)\n    835 def jacfun(*args, **kwargs):\n    836   f_partial, dyn_args = argnums_partial2(fun, argnums, args, kwargs)\n--> 837   tree_map(partial(_check_input_dtype_jacrev, holomorphic, allow_int), dyn_args)\n    838   y, pullback, *maybe_aux = vjp(f_partial, *dyn_args, has_aux=has_aux)\n    839   tree_map(partial(_check_output_dtype_jacrev, holomorphic), y)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/tree_util.py:400, in tree_map(f, tree, is_leaf, *rest)\n    398   err = next(_prefix_error((), tree, r2, is_leaf), None)  # type: ignore\n    399   raise (err('tree_map tree') if err is not None else e) from None\n--> 400 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))\n\nCell In[10], line 6, in <lambda>(_, args)\n----> 6     lambda _, args: MyTree(*args))\n\nCell In[10], line 3, in MyTree.__init__(self, a)\n      2   def __init__(self, a):\n----> 3     self.a = jnp.asarray(a)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/array_constructors.py:456, in asarray(a, dtype, order, copy, device, out_sharding)\n    454 if dtype is not None:\n    455   dtype = dtypes.check_and_canonicalize_user_dtype(dtype, \"asarray\")\n--> 456 return array(a, dtype=dtype, copy=bool(copy), order=order, device=device,\n    457              out_sharding=out_sharding)\n\nFile ~/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/numpy/array_constructors.py:285, in array(object, dtype, copy, order, ndmin, device, out_sharding, *args)\n    282 leaves, treedef = tree_util.tree_flatten(\n    283     object, is_leaf=lambda x: not isinstance(x, (list, tuple)))\n    284 if any(leaf is None for leaf in leaves):\n--> 285   raise ValueError(\"None is not a valid value for jnp.array\")\n    286 leaves = [\n    287     leaf\n    288     if (leaf_jax_array := getattr(leaf, \"__jax_array__\", None)) is None\n    289     else leaf_jax_array()\n    290     for leaf in leaves\n    291 ]\n    292 if dtype is None:\n    293   # Use lattice_result_type rather than result_type to avoid canonicalization.\n    294   # Otherwise, weakly-typed inputs would have their dtypes canonicalized.\n\nValueError: None is not a valid value for jnp.array\n```\n\nExample:\n```text\nclass MyTree:\n  def __init__(self, a):\n    if not (type(a) is object or a is None or isinstance(a, MyTree)):\n      a = jnp.asarray(a)\n    self.a = a\n```\n\nExample:\n```text\ndef tree_unflatten(aux_data, children):\n  del aux_data  # Unused in this class.\n  obj = object.__new__(MyTree)\n  obj.a = children[0]\n  return obj\n```\n\nExample:\n```text\nfrom jax.tree_util import tree_flatten, tree_unflatten\nimport jax.numpy as jnp\n\n# The structured value to be transformed\nvalue_structured = [1., (2., 3.)]\n\n# The leaves in value_flat correspond to the `*` markers in value_tree\nvalue_flat, value_tree = tree_flatten(value_structured)\nprint(f\"{value_flat=}\\n{value_tree=}\")\n\n# Transform the flat value list using an element-wise numeric transformer\ntransformed_flat = list(map(lambda v: v * 2., value_flat))\nprint(f\"{transformed_flat=}\")\n\n# Reconstruct the structured output, using the original\ntransformed_structured = tree_unflatten(value_tree, transformed_flat)\nprint(f\"{transformed_structured=}\")\n```\n\nExample:\n```text\nvalue_flat=[1.0, 2.0, 3.0]\nvalue_tree=PyTreeDef([*, (*, *)])\ntransformed_flat=[2.0, 4.0, 6.0]\ntransformed_structured=[2.0, (4.0, 6.0)]\n```\n\nExample:\n```text\nfrom collections import namedtuple\nPoint = namedtuple('Point', ['x', 'y'])\n\nexample_containers = [\n    (1., [2., 3.]),\n    (1., {'b': 2., 'a': 3.}),\n    1.,\n    None,\n    jnp.zeros(2),\n    Point(1., 2.)\n]\ndef show_example(structured):\n  flat, tree = tree_flatten(structured)\n  unflattened = tree_unflatten(tree, flat)\n  print(f\"{structured=}\\n  {flat=}\\n  {tree=}\\n  {unflattened=}\")\n\nfor structured in example_containers:\n  show_example(structured)\n```\n\nExample:\n```text\nstructured=(1.0, [2.0, 3.0])\n  flat=[1.0, 2.0, 3.0]\n  tree=PyTreeDef((*, [*, *]))\n  unflattened=(1.0, [2.0, 3.0])\nstructured=(1.0, {'b': 2.0, 'a': 3.0})\n  flat=[1.0, 3.0, 2.0]\n  tree=PyTreeDef((*, {'a': *, 'b': *}))\n  unflattened=(1.0, {'a': 3.0, 'b': 2.0})\nstructured=1.0\n  flat=[1.0]\n  tree=PyTreeDef(*)\n  unflattened=1.0\nstructured=None\n  flat=[]\n  tree=PyTreeDef(None)\n  unflattened=None\nstructured=Array([0., 0.], dtype=float32)\n  flat=[Array([0., 0.], dtype=float32)]\n  tree=PyTreeDef(*)\n  unflattened=Array([0., 0.], dtype=float32)\nstructured=Point(x=1.0, y=2.0)\n  flat=[1.0, 2.0]\n  tree=PyTreeDef(CustomNode(namedtuple[Point], [*, *]))\n  unflattened=Point(x=1.0, y=2.0)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.815Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":26,"totalLines":451,"estimatedTokens":3419}}103{"id":"doc-the_checkify_transformation_jax_documentation-3f23e5bb","source":"documentation","title":"The checkify transformation — JAX documentation","url":"https://docs.jax.dev/en/latest/debugging/checkify_guide.html","text":"Example:\n```text\nfrom jax.experimental import checkify\nimport jax\nimport jax.numpy as jnp\n\ndef f(x, i):\n  checkify.check(i >= 0, \"index needs to be non-negative, got {i}\", i=i)\n  y = x[i]\n  z = jnp.sin(y)\n  return z\n\njittable_f = checkify.checkify(f)\n\nerr, z = jax.jit(jittable_f)(jnp.ones((5,)), -2)\nprint(err.get())\n# >> index needs to be non-negative, got -2! (check failed at <...>:6 (f))\n```\n\nExample:\n```text\nerrors = checkify.user_checks | checkify.index_checks | checkify.float_checks\nchecked_f = checkify.checkify(f, errors=errors)\n\nerr, z = checked_f(jnp.ones((5,)), 100)\nerr.throw()\n# ValueError: out-of-bounds indexing at <..>:7 (f)\n\nerr, z = checked_f(jnp.ones((5,)), -1)\nerr.throw()\n# ValueError: index needs to be non-negative! (check failed at <…>:6 (f))\n\nerr, z = checked_f(jnp.array([jnp.inf, 1]), 0)\nerr.throw()\n# ValueError: nan generated by primitive sin at <...>:8 (f)\n\nerr, z = checked_f(jnp.array([5, 1]), 0)\nerr.throw()  # if no error occurred, throw does nothing!\n```\n\nExample:\n```text\njax.jit(f)(jnp.ones((5,)), -1)  # checkify transformation not used\n# ValueError: Cannot abstractly evaluate a checkify.check which was not functionalized.\n```\n\nExample:\n```text\nerr, z = jax.pmap(checked_f)(jnp.ones((3, 5)), jnp.array([-1, 2, 100]))\nerr.throw()\n\"\"\"\nValueError:\n..  at mapped index 0: index needs to be non-negative! (check failed at :6 (f))\n..  at mapped index 2: out-of-bounds indexing at <..>:7 (f)\n\"\"\"\n```\n\nExample:\n```text\ndef f(x):\n  assert x > 0., \"must be positive!\"\n  return jnp.log(x)\n\njax.grad(f)(0.)\n# ValueError: \"must be positive!\"\n```\n\nExample:\n```text\njax.jit(f)(0.)\n# ConcretizationTypeError: \"Abstract tracer value encountered ...\"\n```\n\nExample:\n```text\ndef f_checked(x):\n  error = x <= 0.\n  result = jnp.log(x)\n  return error, result\n\nerr, y = jax.jit(f_checked)(0.)\nif err:\n  raise ValueError(\"must be positive!\")\n# ValueError: \"must be positive!\"\n```\n\nExample:\n```text\ndef f(x):\n  checkify.check(x > 0., \"{} must be positive!\", x)  # convenient but effectful API\n  return jnp.log(x)\n\nf_checked = checkify(f)\n\nerr, x = jax.jit(f_checked)(-1.)\nerr.throw()\n# ValueError: -1. must be positive! (check failed at <...>:2 (f))\n```\n\nExample:\n```text\njnp.arange(3)[5]                # out of bounds\njnp.sin(jnp.inf)                # NaN generated\njnp.ones((5,)) / jnp.arange(5)  # division by zero\n```\n\nExample:\n```text\ndef f(x, i):\n  y = x[i]        # i could be out of bounds.\n  z = jnp.sin(y)  # z could become NaN\n  return z\n\nerrors = checkify.user_checks | checkify.index_checks | checkify.float_checks\nchecked_f = checkify.checkify(f, errors=errors)\n\nerr, z = checked_f(jnp.ones((5,)), 100)\nerr.throw()\n# ValueError: out-of-bounds indexing at <..>:7 (f)\n\nerr, z = checked_f(jnp.array([jnp.inf, 1]), 0)\nerr.throw()\n# ValueError: nan generated by primitive sin at <...>:8 (f)\n```\n\nExample:\n```text\ndef f(x, i):\n  return x[i]\n\ncheckify_of_jit = checkify.checkify(jax.jit(f))\njit_of_checkify = jax.jit(checkify.checkify(f))\nerr, _ =  checkify_of_jit(jnp.ones((5,)), 100)\nerr.get()\n# out-of-bounds indexing at <..>:2 (f)\nerr, _ = jit_of_checkify(jnp.ones((5,)), 100)\n# out-of-bounds indexing at <..>:2 (f)\n```\n\nExample:\n```text\ndef f(x, i):\n  checkify.check(i >= 0, \"index needs to be non-negative!\")\n  return x[i]\n\nchecked_f = checkify.checkify(f, errors=checkify.all_checks)\nerrs, out = jax.vmap(checked_f)(jnp.ones((3, 5)), jnp.array([-1, 2, 100]))\nerrs.throw()\n\"\"\"\nValueError:\n  at mapped index 0: index needs to be non-negative! (check failed at <...>:2 (f))\n  at mapped index 2: out-of-bounds indexing at <...>:3 (f)\n\"\"\"\n```\n\nExample:\n```text\n@jax.vmap\ndef f(x, i):\n  checkify.check(i >= 0, \"index needs to be non-negative!\")\n  return x[i]\n\nchecked_f = checkify.checkify(f, errors=checkify.all_checks)\nerr, out = checked_f(jnp.ones((3, 5)), jnp.array([-1, 2, 100]))\nerr.throw()\n# ValueError: index needs to be non-negative! (check failed at <...>:2 (f))\n```\n\nExample:\n```text\ndef f(x):\n  return x / x\n\nf = checkify.checkify(f, errors=checkify.float_checks)\nf = pjit(\n  f,\n  in_shardings=PartitionSpec('x', None),\n  out_shardings=(None, PartitionSpec('x', None)))\n\nwith jax.sharding.Mesh(mesh.devices, mesh.axis_names):\n err, data = f(input_data)\nerr.throw()\n# ValueError: divided by zero at <...>:4 (f)\n```\n\nExample:\n```text\ndef f(x):\n return x / (1 + jnp.sqrt(x))\n\ngrad_f = jax.grad(f)\n\nerr, _ = checkify.checkify(grad_f, errors=checkify.nan_checks)(0.)\nprint(err.get())\n>> nan generated by primitive mul at <...>:3 (f)\n```\n\nExample:\n```text\n@jax.custom_vjp\ndef assert_gradient_negative(x):\n return x\n\ndef fwd(x):\n return assert_gradient_negative(x), None\n\ndef bwd(_, grad):\n checkify.check(grad < 0, \"gradient needs to be negative!\")\n return (grad,)\n\nassert_gradient_negative.defvjp(fwd, bwd)\n\njax.grad(assert_gradient_negative)(-1.)\n# ValueError: gradient needs to be negative!\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.816Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":16,"totalLines":217,"estimatedTokens":1211}}104{"id":"doc-external_callbacks_jax_documentation-10de1315","source":"documentation","title":"External callbacks — JAX documentation","url":"https://docs.jax.dev/en/latest/external-callbacks.html","text":"Example:\n```text\nimport jax\n\n@jax.jit\ndef f(x):\n  y = x + 1\n  print(\"intermediate value: {}\".format(y))\n  return y * 2\n\nresult = f(2)\n```\n\nExample:\n```text\nintermediate value: JitTracer(~int32[])\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  y = x + 1\n  jax.debug.print(\"intermediate value: {}\", y)\n  return y * 2\n\nresult = f(2)\n```\n\nExample:\n```text\nintermediate value: 3\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\n\ndef f_host(x):\n  # call a numpy (not jax.numpy) operation:\n  return np.sin(x).astype(x.dtype)\n\ndef f(x):\n  result_shape = jax.ShapeDtypeStruct.like(x)\n  return jax.pure_callback(f_host, result_shape, x, vmap_method='sequential')\n\nx = jnp.arange(5.0)\nf(x)\n```\n\nExample:\n```text\nArray([ 0.       ,  0.841471 ,  0.9092974,  0.14112  , -0.7568025],      dtype=float32)\n```\n\nExample:\n```text\njax.jit(f)(x)\n```\n\nExample:\n```text\ndef body_fun(_, x):\n  return _, f(x)\njax.lax.scan(body_fun, None, jnp.arange(5.0))[1]\n```\n\nExample:\n```text\njax.vmap(f)(x)\n```\n\nExample:\n```text\njax.grad(f)(x)\n```\n\nExample:\n```text\nValueError: Pure callbacks do not support JVP. Please use `jax.custom_jvp` to use callbacks while taking gradients.\n```\n\nExample:\n```text\ndef print_something():\n  print('printing something')\n  return np.int32(0)\n\n@jax.jit\ndef f1():\n  return jax.pure_callback(print_something, np.int32(0))\nf1();\n```\n\nExample:\n```text\nprinting something\n```\n\nExample:\n```text\n@jax.jit\ndef f2():\n  jax.pure_callback(print_something, np.int32(0))\n  return 1.0\nf2();\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\n\ndef raise_via_callback(x):\n  def _raise(x):\n    raise ValueError(f\"value of x is {x}\")\n  return jax.pure_callback(_raise, x, x)\n\ndef raise_if_negative(x):\n  return jax.lax.cond(x < 0, raise_via_callback, lambda x: x, x)\n\nx_batch = jnp.arange(4)\n\n[raise_if_negative(x) for x in x_batch]  # does not raise\n\njax.vmap(raise_if_negative)(x_batch)  # ValueError: value of x is 0\n```\n\nExample:\n```text\nfrom jax.experimental import io_callback\nfrom functools import partial\n\nglobal_rng = np.random.default_rng(0)\n\ndef host_side_random_like(x):\n  \"\"\"Generate a random array like x using the global_rng state\"\"\"\n  # We have two side-effects here:\n  # - printing the shape and dtype\n  # - calling global_rng, thus updating its state\n  print(f'generating {x.dtype}{list(x.shape)}')\n  return global_rng.uniform(size=x.shape).astype(x.dtype)\n\n@jax.jit\ndef numpy_random_like(x):\n  return io_callback(host_side_random_like, x, x)\n\nx = jnp.zeros(5)\nnumpy_random_like(x)\n```\n\nExample:\n```text\ngenerating float32[5]\n```\n\nExample:\n```text\nArray([0.6369617 , 0.26978672, 0.04097353, 0.01652764, 0.8132702 ],      dtype=float32)\n```\n\nExample:\n```text\njax.vmap(numpy_random_like)(x)\n```\n\nExample:\n```text\ngenerating float32[]\ngenerating float32[]\ngenerating float32[]\ngenerating float32[]\ngenerating float32[]\n```\n\nExample:\n```text\nArray([0.91275555, 0.60663575, 0.72949654, 0.543625  , 0.9350724 ],      dtype=float32)\n```\n\nExample:\n```text\n@jax.jit\ndef numpy_random_like_ordered(x):\n  return io_callback(host_side_random_like, x, x, ordered=True)\n\njax.vmap(numpy_random_like_ordered)(x)\n```\n\nExample:\n```text\nValueError: Cannot `vmap` ordered IO callback.\n```\n\nExample:\n```text\ndef body_fun(_, x):\n  return _, numpy_random_like_ordered(x)\njax.lax.scan(body_fun, None, jnp.arange(5.0))[1]\n```\n\nExample:\n```text\nArray([0.81585354, 0.0027385 , 0.8574043 , 0.03358557, 0.72965544],      dtype=float32)\n```\n\nExample:\n```text\njax.grad(numpy_random_like)(x)\n```\n\nExample:\n```text\nValueError: IO callbacks do not support JVP.\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  io_callback(lambda: print('hello'), None)\n  return x\n\njax.grad(f)(1.0);\n```\n\nExample:\n```text\nhello\n```\n\nExample:\n```text\nfrom jax import debug\n\ndef log_value(x):\n  # This could be an actual logging call; we'll use\n  # print() for demonstration\n  print(\"log:\", x)\n\n@jax.jit\ndef f(x):\n  debug.callback(log_value, x)\n  return x\n\nf(1.0);\n```\n\nExample:\n```text\nlog: 1.0\n```\n\nExample:\n```text\nx = jnp.arange(5.0)\njax.vmap(f)(x);\n```\n\nExample:\n```text\nlog: 0.0\nlog: 1.0\nlog: 2.0\nlog: 3.0\nlog: 4.0\n```\n\nExample:\n```text\njax.grad(f)(1.0);\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nimport scipy.special\n\ndef jv(v, z):\n  v, z = jnp.asarray(v), jnp.asarray(z)\n\n  # Require the order v to be integer type: this simplifies\n  # the JVP rule below.\n  assert jnp.issubdtype(v.dtype, jnp.integer)\n\n  # Promote the input to inexact (float/complex).\n  # Note that jnp.result_type() accounts for the enable_x64 flag.\n  z = z.astype(jnp.result_type(float, z.dtype))\n\n  # Wrap scipy function to return the expected dtype.\n  _scipy_jv = lambda v, z: scipy.special.jv(v, z).astype(z.dtype)\n\n  # Define the expected shape & dtype of output.\n  result_shape_dtype = jax.ShapeDtypeStruct(\n      shape=jnp.broadcast_shapes(v.shape, z.shape),\n      dtype=z.dtype)\n\n  # Use vmap_method=\"broadcast_all\" because scipy.special.jv handles broadcasted inputs.\n  return jax.pure_callback(_scipy_jv, result_shape_dtype, v, z, vmap_method=\"broadcast_all\")\n```\n\nExample:\n```text\nfrom functools import partial\nj1 = partial(jv, 1)\nz = jnp.arange(5.0)\n```\n\nExample:\n```text\nprint(j1(z))\n```\n\nExample:\n```text\n[ 0.          0.44005057  0.5767248   0.33905897 -0.06604332]\n```\n\nExample:\n```text\nprint(jax.jit(j1)(z))\n```\n\nExample:\n```text\nprint(jax.vmap(j1)(z))\n```\n\nExample:\n```text\njax.grad(j1)(z)\n```\n\nExample:\n```text\njv = jax.custom_jvp(jv)\n\n@jv.defjvp\ndef _jv_jvp(primals, tangents):\n  v, z = primals\n  _, z_dot = tangents  # Note: v_dot is always 0 because v is integer.\n  jv_minus_1, jv_plus_1 = jv(v - 1, z), jv(v + 1, z)\n  djv_dz = jnp.where(v == 0, -jv_plus_1, 0.5 * (jv_minus_1 - jv_plus_1))\n  return jv(v, z), z_dot * djv_dz\n```\n\nExample:\n```text\nj1 = partial(jv, 1)\nprint(jax.grad(j1)(2.0))\n```\n\nExample:\n```text\n-0.06447162\n```\n\nExample:\n```text\njax.hessian(j1)(2.0)\n```\n\nExample:\n```text\nArray(-0.4003078, dtype=float32, weak_type=True)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.817Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":46,"totalLines":371,"estimatedTokens":1484}}105{"id":"doc-ref_mutable_arrays_for_data_plumbing_and_memory_-e54519f6","source":"documentation","title":"Ref: mutable arrays for data plumbing and memory control — JAX documentation","url":"https://docs.jax.dev/en/latest/array_refs.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n\nx_ref = jax.new_ref(jnp.zeros(3))  # new array ref, with initial value [0., 0., 0.]\n\n@jax.jit\ndef f():\n  x_ref[1] += 1.  # indexed add-update\n\nprint(x_ref)  # Ref([0., 0., 0.])\nf()\nf()\nprint(x_ref)  # Ref([0., 2., 0.])\n```\n\nExample:\n```text\nRef([0., 0., 0.], dtype=float32)\nRef([0., 2., 0.], dtype=float32)\n```\n\nExample:\n```text\ndef g(x):\n  x_ref = jax.new_ref(0.)\n  x_ref[...] = jnp.sin(x)\n  return x_ref[...]\n\nprint(jax.grad(g)(1.0))  # 0.54\n```\n\nExample:\n```text\n0.5403023\n```\n\nExample:\n```text\nx_ref = jax.new_ref(1.0)\ntry:\n  jnp.sin(x_ref)  # error! can't do math on refs\nexcept Exception as e:\n  print(e)\n```\n\nExample:\n```text\nsin requires ndarray or scalar arguments, got <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'> at position 0.\n```\n\nExample:\n```text\nfrom jax import Array, Ref\n\ndef array_ref(init_val: Array) -> Ref:\n  \"\"\"Introduce a new reference with given initial value.\"\"\"\n```\n\nExample:\n```text\ndef freeze(ref: Ref) -> Array:\n  \"\"\"Invalidate given reference and produce its final value.\"\"\"\n```\n\nExample:\n```text\nimport types\nIndex = int | slice | Array | types.EllipsisType\nIndexer = Index | tuple[Index, ...]\n\ndef get(ref: Ref, idx: Indexer) -> Array:\n  \"\"\"Returns `ref[idx]` for NumPy-style indexer `idx`.\"\"\"\n\ndef swap(ref: Ref, idx: Indexer, val: Array) -> Array:\n  \"\"\"Performs `newval, ref[idx] = ref[idx], val` and returns `newval`.\"\"\"\n```\n\nExample:\n```text\nx_ref = jax.new_ref(jnp.arange(12.).reshape(3, 4))\n\n# int indexing\nrow = x_ref[0]\nx_ref[1] = row\n\n# tuple indexing\nval = x_ref[1, 2]\nx_ref[2, 3] = val\n\n# slice indexing\ncol = x_ref[:, 1]\nx_ref[0, :3] = col\n\n# advanced int array indexing\nvals = x_ref[jnp.array([0, 0, 1]), jnp.array([1, 2, 3])]\nx_ref[jnp.array([1, 2, 1]), jnp.array([0, 0, 1])] = vals\n```\n\nExample:\n```text\n# takes ref as an argument => impure\n@jax.jit\ndef impure1(x_ref, y_ref):\n  x_ref[...] = y_ref[...]\n\n# closes over ref => impure\ny_ref = jax.new_ref(0)\n\n@jax.jit\ndef impure2(x):\n  y_ref[...] = x\n```\n\nExample:\n```text\n# internal refs => still pure\n@jax.jit\ndef pure1(x):\n  ref = jax.new_ref(x)\n  ref[...] = ref[...] + ref[...]\n  return ref[...]\n```\n\nExample:\n```text\nx_ref = jax.new_ref(0.)\n\n# can't return refs\n@jax.jit\ndef err1(x_ref):\n  x_ref[...] = 5.\n  return x_ref  # error!\ntry:\n  err1(x_ref)\nexcept Exception as e:\n  print(e)\n\n# can't pass a ref as an argument more than once\n@jax.jit\ndef err2(x_ref, y_ref):\n  ...\ntry:\n  err2(x_ref, x_ref)  # error!\nexcept Exception as e:\n  print(e)\n\n# can't pass and close over the same ref\n@jax.jit\ndef err3(y_ref):\n  y_ref[...] = x_ref[...]\ntry:\n  err3(x_ref)  # error!\nexcept Exception as e:\n  print(e)\n\n# can only freeze in creation scope\n@jax.jit\ndef err4(x_ref):\n  jax.freeze(x_ref)\ntry:\n  err4(x_ref)  # error!\nexcept Exception as e:\n  print(e)\n```\n\nExample:\n```text\nfunction err1 at /tmp/ipykernel_1344/3915325362.py:4 traced for jit returned a mutable array reference of type Ref{float32[]}, but mutable array references cannot be returned.\n\nThe returned mutable array was passed in as the argument x_ref.\nonly one reference to a mutable array may be passed as an argument to a function, but when tracing err2 at /tmp/ipykernel_1344/3915325362.py:14 for jit the mutable array reference of type Ref{float32[]} appeared at both x_ref and y_ref.\nwhen tracing err3 at /tmp/ipykernel_1344/3915325362.py:23 for jit, a mutable array reference of type Ref{float32[]} was both closed over and passed as the argument y_ref\n```\n\nExample:\n```text\n# vmap over ref args is okay\ndef dist(x, y, out_ref):\n  assert x.ndim == y.ndim == 1\n  assert out_ref.ndim == 0\n  out_ref[...] = jnp.sum((x - y) ** 2)\n\nvecs = jnp.arange(12.).reshape(3, 4)\nout_ref = jax.new_ref(jnp.zeros((3, 3)))\njax.vmap(jax.vmap(dist, (0, None, 0)), (None, 0, 0))(vecs, vecs, out_ref)  # ok!\nprint(out_ref)\n```\n\nExample:\n```text\nRef([[  0.,  64., 256.],\n       [ 64.,   0.,  64.],\n       [256.,  64.,   0.]], dtype=float32)\n```\n\nExample:\n```text\n# vmap with a closed-over ref is not\nx_ref = jax.new_ref(0.)\n\ndef err5(x):\n  x_ref[...] = x\n\ntry:\n  jax.vmap(err5)(jnp.arange(3.))  # error!\nexcept Exception as e:\n  print(e)\n```\n\nExample:\n```text\nperforming a set/swap operation with vmapped value on an unbatched array reference of type Ref{float32[]}. Move the array reference to be an argument to the vmapped function?\n```\n\nExample:\n```text\n@jax.jit\ndef pure2(x):\n  ref = jax.new_ref(x)\n  ref[...] = ref[...] + ref[...]\n  return ref[...]\n\nprint(jax.grad(pure2)(3.0))  # 2.0\n```\n\nExample:\n```text\n2.0\n```\n\nExample:\n```text\n# error\ndef err6(x, some_plumbing_ref):\n  y = x + x\n  some_plumbing_ref[...] += y\n  return y\n\n# fine\ndef foo(x, some_plumbing_ref):\n  y = x + x\n  some_plumbing_ref[...] += jax.lax.stop_gradient(y)\n  return y\n```\n\nExample:\n```text\n# First, define the helper `stash_grads`:\n\n@jax.custom_vjp\ndef stash_grads(grads_ref, x):\n  return x\n\ndef stash_grads_fwd(grads_ref, x):\n  return x, grads_ref\n\ndef stash_grads_bwd(grads_ref, g):\n  grads_ref[...] = g\n  return None, g\n\nstash_grads.defvjp(stash_grads_fwd, stash_grads_bwd)\n```\n\nExample:\n```text\n# Now, use `stash_grads` to stash intermediate gradients:\n\ndef f(x, grads_ref):\n  x = stash_grads(grads_ref, x)\n  x = jnp.sin(x)\n  return x\n\ngrads_ref = jax.new_ref(0.)\njax.grad(f)(1., grads_ref)\nprint(grads_ref)  # Ref(0.54), the gradient at the stash point: cos(1.)\n```\n\nExample:\n```text\nRef(0.5403023, dtype=float32)\n```\n\nExample:\n```text\ndef f(x_ref):\n  return x_ref[...] ** 2\n\nx_ref = jax.new_ref(2.)\ny, f_vjp = jax.vjp(f, x_ref)\n\nx_grad_ref = jax.new_ref(0.)\nf_vjp.with_refs(x_grad_ref)(1.0)  # bind the gradient ref, then apply the VJP\nprint(x_grad_ref)  # Ref(4.)\n```\n\nExample:\n```text\nRef(4., dtype=float32, weak_type=True)\n```\n\nExample:\n```text\n_, f_vjp = jax.vjp(f, jax.new_ref(2.))\ntry:\n  f_vjp(1.0)  # error! no ref bound for the ref-typed argument's gradient\nexcept Exception as e:\n  print(e)  # ... gradient must be accumulated into a ref ... `with_refs` ...\n```\n\nExample:\n```text\nthe argument at position args[0] of the differentiated function f at /tmp/ipykernel_1344/324718511.py:1 is Ref-typed, so its gradient must be accumulated into a ref, but no gradient ref was provided. Bind one using the VJP function's `with_refs` method before applying it, as in `f_vjp.with_refs(grad_ref)(ct)`; the gradient will be accumulated into `grad_ref` in-place via addition. Or, to skip computing this argument's gradient, pass `jax.ad.DontWant()` in place of a gradient ref.\n```\n\nExample:\n```text\nx_grad_ref = jax.new_ref(100.)\n_, f_vjp = jax.vjp(f, jax.new_ref(2.))\nf_vjp.with_refs(x_grad_ref)(1.0)\nprint(x_grad_ref)  # Ref(104.), i.e. 100. + 4.: accumulated, not set\n```\n\nExample:\n```text\nRef(104., dtype=float32, weak_type=True)\n```\n\nExample:\n```text\ndef g(x_ref):\n  x_ref[...] = jnp.sin(x_ref[...])\n  return x_ref[...] ** 2\n\nx_ref = jax.new_ref(2.)\n_, g_vjp = jax.vjp(g, x_ref)  # runs g, so x_ref is updated in-place here\ng_grad_ref = jax.new_ref(0.)\ng_vjp.with_refs(g_grad_ref)(1.0)\nprint(g_grad_ref)  # Ref(-0.757), i.e. 2*sin(2)*cos(2)\n```\n\nExample:\n```text\nRef(-0.7568025, dtype=float32, weak_type=True)\n```\n\nExample:\n```text\n_, sin_vjp = jax.vjp(jnp.sin, 1.0)\nx_bar, = sin_vjp(1.0)  # the usual way: gradient returned as a value\nprint(x_bar)  # 0.54\n\ngrad_ref = jax.new_ref(0.)\n_, sin_vjp = jax.vjp(jnp.sin, 1.0)\nsin_vjp.with_refs(grad_ref)(1.0)  # gradient accumulated into grad_ref\nprint(grad_ref)  # Ref(0.54)\n```\n\nExample:\n```text\n0.5403023\nRef(0.5403023, dtype=float32, weak_type=True)\n```\n\nExample:\n```text\n@jax.jit\ndef take(x, i):\n  return x[i]\n\nx = jnp.arange(10.)\n\n_, take_vjp = jax.vjp(take, x, 3)\nx_bar, _ = take_vjp(1.0)\nprint(x_bar)  # [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.]\n```\n\nExample:\n```text\n[0. 0. 0. 1. 0. 0. 0. 0. 0. 0.]\n```\n\nExample:\n```text\ngrad_ref = jax.new_ref(jnp.zeros(10))\n\nfor i in [3, 5, 3]:\n  _, take_vjp = jax.vjp(take, x, i)\n  take_vjp.with_refs(grad_ref, jax.ad.GradValue())(1.0)  # no gradient ref for i\n\nprint(grad_ref)  # Ref([0., 0., 0., 2., 0., 1., 0., 0., 0., 0.])\n```\n\nExample:\n```text\nRef([0., 0., 0., 2., 0., 1., 0., 0., 0., 0.], dtype=float32)\n```\n\nExample:\n```text\n@jax.make_jaxpr\ndef take_vjp_jaxpr():\n  _, take_vjp = jax.vjp(take, x, 3)\n  take_vjp.with_refs(grad_ref, jax.ad.GradValue())(1.0)\n\nprint(take_vjp_jaxpr())\n```\n\nExample:\n```text\n{ lambda a:f32[10] b:Ref{f32[10]}; . let\n    _:f32[] c:i32[] = jit[\n      name=take\n      jaxpr={ lambda ; a:f32[10] d:i32[]. let\n          e:bool[] = lt d 0:i32[]\n          f:i32[] = convert_element_type[new_dtype=int32 weak_type=False] d\n          g:i32[] = add f 10:i32[]\n          c:i32[] = select_n e d g\n          h:f32[1] = dynamic_slice[slice_sizes=(1,)] a c\n          _:f32[] = squeeze[dimensions=(0,)] h\n        in (_, c) }\n    ] a 3:i32[]\n    jit[\n      name=take\n      jaxpr={ lambda ; c:i32[] b:Ref{f32[10]} i:f32[]. let\n          j:f32[1] = broadcast_in_dim i\n          b[c:c+1] += j\n        in () }\n    ] c b 1.0:f32[]\n  in () }\n```\n\nExample:\n```text\ndef predict(W, x):\n  return W @ x\n\nW = jnp.ones((4, 4))\nx = jnp.ones(4)\n\n_, f_vjp = jax.vjp(predict, W, x)\nW_bar, x_bar = f_vjp.with_refs(jax.ad.GradValue(), jax.ad.DontWant())(jnp.ones(4))\nprint(W_bar[0])  # [1., 1., 1., 1.]\nprint(x_bar)  # DidntWant()\n```\n\nExample:\n```text\n[1. 1. 1. 1.]\nDidntWant()\n```\n\nExample:\n```text\n_, f_vjp = jax.vjp(predict, W, x)\nboth = jax.make_jaxpr(lambda: f_vjp(jnp.ones(4)))()\nonly_W = jax.make_jaxpr(\n    lambda: f_vjp.with_refs(jax.ad.GradValue(), jax.ad.DontWant())(jnp.ones(4)))()\n\nprint(str(both).count('dot_general'))    # 2, one dot for each gradient\nprint(str(only_W).count('dot_general'))  # 1, the dot for x_bar is skipped\n```\n\nExample:\n```text\n2\n1\n```\n\nExample:\n```text\nNUM_LAYERS = 3\nNUM_MUBATCHES = 5\nMUBATCH_SIZE = 7\n\ndef mubatch_loss(Ws, xs):\n  # inner loop over layers\n  act, _ = jax.lax.scan(lambda x, W: (jnp.dot(x, W), None), xs, Ws)\n  return jnp.mean(act)\n\ndef process_batch(Ws, xs_batch):\n  grad_acc = jax.new_ref(jnp.zeros_like(Ws))\n\n  def process_mubatch(_, xs):\n    loss, f_vjp = jax.vjp(lambda Ws: mubatch_loss(Ws, xs), Ws)\n    f_vjp.with_refs(grad_acc)(jnp.ones_like(loss))  # accumulate in-place\n    return (), loss\n\n  xs_mubatches = xs_batch.reshape(NUM_MUBATCHES, MUBATCH_SIZE, -1)\n  # outer loop over microbatches\n  (), losses = jax.lax.scan(process_mubatch, (), xs_mubatches)\n  return jax.freeze(grad_acc), losses\n\nWs = jnp.ones((NUM_LAYERS, 4, 4))\nxs_batch = jnp.ones((NUM_MUBATCHES * MUBATCH_SIZE, 4))\ngrads, losses = process_batch(Ws, xs_batch)\n```\n\nExample:\n```text\nxs_mubatches = xs_batch.reshape(NUM_MUBATCHES, MUBATCH_SIZE, -1)\ngrads_expected = jax.grad(\n    lambda Ws: sum(mubatch_loss(Ws, xs) for xs in xs_mubatches))(Ws)\nprint(jnp.allclose(grads, grads_expected, atol=1e-3, rtol=1e-3))  # True\n```\n\nExample:\n```text\nTrue\n```\n\nExample:\n```text\n@jax.jit\ndef sin_inplace(x_ref):\n  x_ref[...] = jnp.sin(x_ref[...])\n\nx_ref = jax.new_ref(jnp.arange(3.))\nprint(x_ref.unsafe_buffer_pointer(), x_ref)\nsin_inplace(x_ref)\nprint(x_ref.unsafe_buffer_pointer(), x_ref)\n```\n\nExample:\n```text\n103907990956608 Ref([0., 1., 2.], dtype=float32)\n103907990956608 Ref([0.        , 0.84147096, 0.9092974 ], dtype=float32)\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax.lax import scan\n\ndef foreach(*args):\n  def decorator(body):\n    return scan(lambda _, elts: (None, body(*elts)), None, args)[1]\n  return decorator\n```\n\nExample:\n```text\nr = jax.new_ref(0)\nxs = jnp.arange(10)\n\n@foreach(xs)\ndef ys(x):\n  r[...] += x\n  return x * 2\n\nprint(r)   # Ref(45, dtype=int32)\nprint(ys)  # [ 0  2  4  6  8 10 12 14 16 18]\n```\n\nExample:\n```text\nRef(45, dtype=int32)\n[ 0  2  4  6  8 10 12 14 16 18]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.819Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":52,"totalLines":563,"estimatedTokens":2909}}106{"id":"doc-xla_compiler_flags_jax_documentation-6f902b47","source":"documentation","title":"XLA compiler flags — JAX documentation","url":"https://docs.jax.dev/en/latest/xla_flags.html","text":"Example:\n```text\nimport os\n\n# Set multiple flags separated by spaces\nos.environ['XLA_FLAGS'] = '--flag1=value1 --flag2=value2'\n```\n\nExample:\n```text\nXLA_FLAGS='--flag1=value1 --flag2=value2'  python3 source.py\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.819Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":2,"totalLines":14,"estimatedTokens":57}}107{"id":"doc-ahead_of_time_lowering_and_compilation_jax_docum-01fc4991","source":"documentation","title":"Ahead-of-time lowering and compilation — JAX documentation","url":"https://docs.jax.dev/en/latest/aot.html","text":"Example:\n```text\n>>> import jax\n\n>>> def f(x, y): return 2 * x + y\n>>> x, y = 3, 4\n\n>>> traced = jax.jit(f).trace(x, y)\n\n>>> # Print the specialized, staged-out representation (as Jaxpr IR)\n>>> print(traced.jaxpr)\n{ lambda ; a:i32[] b:i32[]. let\n    c:i32[] = mul 2:i32[] a\n    d:i32[] = add c b\n  in (d,) }\n\n>>> lowered = traced.lower()\n\n>>> # Print lowered HLO\n>>> print(lowered.as_text())\nmodule @jit_f attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {\n  func.func public @main(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32> {jax.result_info = \"result\"}) {\n    %c = stablehlo.constant dense<2> : tensor<i32>\n    %0 = stablehlo.multiply %c, %arg0 : tensor<i32>\n    %1 = stablehlo.add %0, %arg1 : tensor<i32>\n    return %1 : tensor<i32>\n  }\n}\n\n>>> compiled = lowered.compile()\n\n>>> # Query for cost analysis, print FLOP estimate\n>>> compiled.cost_analysis()['flops']\n2.0\n\n>>> # Execute the compiled function!\n>>> compiled(x, y)\nArray(10, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\n>>> i32_scalar = jax.ShapeDtypeStruct((), jnp.dtype('int32'))\n>>> jax.jit(f).trace(i32_scalar, i32_scalar).lower().compile()(x, y)\nArray(10, dtype=int32)\n```\n\nExample:\n```text\n>>> x_1d = y_1d = jnp.arange(3)\n>>> jax.jit(f).trace(i32_scalar, i32_scalar).lower().compile()(x_1d, y_1d)  \n...\nTraceback (most recent call last):\nTypeError: Argument types differ from the types for which this computation was compiled. The mismatches are:\nArgument 'x' compiled with int32[] and called with int32[3]\nArgument 'y' compiled with int32[] and called with int32[3]\n\n>>> x_f = y_f = jnp.float32(72.)\n>>> jax.jit(f).trace(i32_scalar, i32_scalar).lower().compile()(x_f, y_f)  \n...\nTraceback (most recent call last):\nTypeError: Argument types differ from the types for which this computation was compiled. The mismatches are:\nArgument 'x' compiled with int32[] and called with float32[]\nArgument 'y' compiled with int32[] and called with float32[]\n```\n\nExample:\n```text\n>>> lowered_with_x = jax.jit(f, static_argnums=0).trace(7, 8).lower()\n\n>>> # Lowered HLO, specialized to the *value* of the first argument (7)\n>>> print(lowered_with_x.as_text())\nmodule @jit_f attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {\n  func.func public @main(%arg0: tensor<i32>) -> (tensor<i32> {jax.result_info = \"result\"}) {\n    %c = stablehlo.constant dense<14> : tensor<i32>\n    %0 = stablehlo.add %c, %arg0 : tensor<i32>\n    return %0 : tensor<i32>\n  }\n}\n\n>>> lowered_with_x.compile()(5)\nArray(19, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\n>>> jax.jit(f, static_argnums=0).trace(i32_scalar, i32_scalar)  \nTraceback (most recent call last):\nTypeError: unsupported operand type(s) for *: 'int' and 'ShapeDtypeStruct'\n\n>>> jax.jit(f, static_argnums=0).trace(10, i32_scalar).lower().compile()(5)\nArray(25, dtype=int32)\n```\n\nExample:\n```text\n>>> def g(x):\n...   assert x.shape == (3, 2)\n...   return x @ jnp.ones(2)\n\n>>> def make_z(*shape):\n...   return jnp.arange(np.prod(shape)).reshape(shape)\n\n>>> z, zs = make_z(3, 2), make_z(4, 3, 2)\n\n>>> g_jit = jax.jit(g)\n>>> g_aot = jax.jit(g).trace(z).lower().compile()\n\n>>> jax.vmap(g_jit)(zs)\nArray([[ 1.,  5.,  9.],\n       [13., 17., 21.],\n       [25., 29., 33.],\n       [37., 41., 45.]], dtype=float32)\n\n>>> jax.vmap(g_aot)(zs)  \nTraceback (most recent call last):\nTypeError: Cannot apply JAX transformations to a function lowered and compiled for a particular signature. Detected argument of Tracer type <class 'jax._src.interpreters.batching.BatchTracer'>\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.820Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":6,"totalLines":120,"estimatedTokens":886}}108{"id":"doc-autobatching_for_bayesian_inference_jax_document-8cd70e87","source":"documentation","title":"Autobatching for Bayesian inference — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/vmapped_log_probs.html","text":"Example:\n```text\nimport matplotlib.pyplot as plt\n\nimport jax\n\nimport jax.numpy as jnp\nimport jax.scipy as jsp\nfrom jax import random\n\nimport numpy as np\nimport scipy as sp\n```\n\nExample:\n```text\nnp.random.seed(10009)\n\nnum_features = 10\nnum_points = 100\n\ntrue_beta = np.random.randn(num_features).astype(jnp.float32)\nall_x = np.random.randn(num_points, num_features).astype(jnp.float32)\ny = (np.random.rand(num_points) < sp.special.expit(all_x.dot(true_beta))).astype(jnp.int32)\n```\n\nExample:\n```text\ny\n```\n\nExample:\n```text\narray([0, 0, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 0, 0, 0, 1, 1, 1, 1, 0,\n       1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0,\n       1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0,\n       0, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 1,\n       1, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0], dtype=int32)\n```\n\nExample:\n```text\ndef log_joint(beta):\n    result = 0.\n    # Note that no `axis` parameter is provided to `jnp.sum`.\n    result = result + jnp.sum(jsp.stats.norm.logpdf(beta, loc=0., scale=1.))\n    result = result + jnp.sum(-jnp.log(1 + jnp.exp(-(2*y-1) * jnp.dot(all_x, beta))))\n    return result\n```\n\nExample:\n```text\nlog_joint(np.random.randn(num_features))\n```\n\nExample:\n```text\nArray(-213.23558, dtype=float32)\n```\n\nExample:\n```text\n# This doesn't work, because we didn't write `log_prob()` to handle batching.\ntry:\n  batch_size = 10\n  batched_test_beta = np.random.randn(batch_size, num_features)\n\n  log_joint(np.random.randn(batch_size, num_features))\nexcept ValueError as e:\n  print(\"Caught expected exception \" + str(e))\n```\n\nExample:\n```text\nCaught expected exception Incompatible shapes for broadcasting: shapes=[(100,), (100, 10)]\n```\n\nExample:\n```text\ndef batched_log_joint(beta):\n    result = 0.\n    # Here (and below) `sum` needs an `axis` parameter. At best, forgetting to set axis\n    # or setting it incorrectly yields an error; at worst, it silently changes the\n    # semantics of the model.\n    result = result + jnp.sum(jsp.stats.norm.logpdf(beta, loc=0., scale=1.),\n                           axis=-1)\n    # Note the multiple transposes. Getting this right is not rocket science,\n    # but it's also not totally mindless. (I didn't get it right on the first\n    # try.)\n    result = result + jnp.sum(-jnp.log(1 + jnp.exp(-(2*y-1) * jnp.dot(all_x, beta.T).T)),\n                           axis=-1)\n    return result\n```\n\nExample:\n```text\nbatch_size = 10\nbatched_test_beta = np.random.randn(batch_size, num_features)\n\nbatched_log_joint(batched_test_beta)\n```\n\nExample:\n```text\nArray([-147.84032 , -207.02205 , -109.26076 , -243.80832 , -163.02908 ,\n       -143.84848 , -160.28772 , -113.7717  , -126.605446, -190.81989 ],      dtype=float32)\n```\n\nExample:\n```text\nvmap_batched_log_joint = jax.vmap(log_joint)\nvmap_batched_log_joint(batched_test_beta)\n```\n\nExample:\n```text\n@jax.jit\ndef log_joint(beta):\n    result = 0.\n    # Note that no `axis` parameter is provided to `jnp.sum`.\n    result = result + jnp.sum(jsp.stats.norm.logpdf(beta, loc=0., scale=10.))\n    result = result + jnp.sum(-jnp.log(1 + jnp.exp(-(2*y-1) * jnp.dot(all_x, beta))))\n    return result\n\nbatched_log_joint = jax.jit(jax.vmap(log_joint))\n```\n\nExample:\n```text\ndef elbo(beta_loc, beta_log_scale, epsilon):\n    beta_sample = beta_loc + jnp.exp(beta_log_scale) * epsilon\n    return jnp.mean(batched_log_joint(beta_sample), 0) + jnp.sum(beta_log_scale - 0.5 * np.log(2*np.pi))\n\nelbo = jax.jit(elbo)\nelbo_val_and_grad = jax.jit(jax.value_and_grad(elbo, argnums=(0, 1)))\n```\n\nExample:\n```text\ndef normal_sample(key, shape):\n    \"\"\"Convenience function for quasi-stateful RNG.\"\"\"\n    new_key, sub_key = random.split(key)\n    return new_key, random.normal(sub_key, shape)\n\nnormal_sample = jax.jit(normal_sample, static_argnums=(1,))\n\nkey = random.key(10003)\n\nbeta_loc = jnp.zeros(num_features, jnp.float32)\nbeta_log_scale = jnp.zeros(num_features, jnp.float32)\n\nstep_size = 0.01\nbatch_size = 128\nepsilon_shape = (batch_size, num_features)\nfor i in range(1000):\n    key, epsilon = normal_sample(key, epsilon_shape)\n    elbo_val, (beta_loc_grad, beta_log_scale_grad) = elbo_val_and_grad(\n        beta_loc, beta_log_scale, epsilon)\n    beta_loc += step_size * beta_loc_grad\n    beta_log_scale += step_size * beta_log_scale_grad\n    if i % 10 == 0:\n        print('{}\\t{}'.format(i, elbo_val))\n```\n\nExample:\n```text\n0\t-175.5615997314453\n10\t-112.76364135742188\n20\t-102.41358184814453\n30\t-100.27793884277344\n40\t-99.55818176269531\n50\t-98.18000793457031\n60\t-98.60237884521484\n70\t-97.69735717773438\n80\t-97.53225708007812\n90\t-97.17940521240234\n100\t-97.09412384033203\n110\t-97.4031753540039\n120\t-97.04466247558594\n130\t-97.20584106445312\n140\t-96.89036560058594\n150\t-96.91874694824219\n160\t-97.00558471679688\n170\t-97.45591735839844\n180\t-96.73572540283203\n190\t-96.95585632324219\n200\t-97.51351928710938\n210\t-96.92330932617188\n220\t-97.03160095214844\n230\t-96.88632202148438\n240\t-96.9697036743164\n250\t-97.35342407226562\n260\t-97.07598876953125\n270\t-97.24360656738281\n280\t-97.23467254638672\n290\t-97.02444458007812\n300\t-97.00311279296875\n310\t-97.07694244384766\n320\t-97.33139038085938\n330\t-97.15113830566406\n340\t-97.28958129882812\n350\t-97.41972351074219\n360\t-96.95799255371094\n370\t-97.36982727050781\n380\t-97.00273132324219\n390\t-97.10066223144531\n400\t-97.13653564453125\n410\t-96.87237548828125\n420\t-97.24083709716797\n430\t-97.04019165039062\n440\t-96.68864440917969\n450\t-97.19795989990234\n460\t-97.18959045410156\n470\t-97.09814453125\n480\t-97.11341857910156\n490\t-97.20771789550781\n500\t-97.39350128173828\n510\t-97.25328063964844\n520\t-97.20198822021484\n530\t-96.95065307617188\n540\t-97.37590789794922\n550\t-96.98526763916016\n560\t-97.0145263671875\n570\t-96.9732894897461\n580\t-97.04313659667969\n590\t-97.38460540771484\n600\t-97.31581115722656\n610\t-97.10185241699219\n620\t-97.22990417480469\n630\t-97.18515014648438\n640\t-97.15637969970703\n650\t-97.13624572753906\n660\t-97.0641860961914\n670\t-97.17774200439453\n680\t-97.31779479980469\n690\t-97.42807006835938\n700\t-97.18154907226562\n710\t-97.57279968261719\n720\t-96.99563598632812\n730\t-97.15852355957031\n740\t-96.85628509521484\n750\t-96.8902587890625\n760\t-97.11228942871094\n770\t-97.214111328125\n780\t-96.99479675292969\n790\t-97.30390930175781\n800\t-96.98690795898438\n810\t-97.12834167480469\n820\t-97.51512145996094\n830\t-97.41466522216797\n840\t-96.89874267578125\n850\t-96.84567260742188\n860\t-97.2318344116211\n870\t-97.24137115478516\n880\t-96.74853515625\n890\t-97.09489440917969\n900\t-97.13866424560547\n910\t-96.79051971435547\n920\t-97.06621551513672\n930\t-97.14911651611328\n940\t-97.26902770996094\n950\t-97.01964569091797\n960\t-96.95348358154297\n970\t-97.138916015625\n980\t-97.60130310058594\n990\t-97.25077056884766\n```\n\nExample:\n```text\nplt.figure(figsize=(7, 7))\nplt.plot(true_beta, beta_loc, '.', label='Approximated Posterior Means')\nplt.plot(true_beta, beta_loc + 2*jnp.exp(beta_log_scale), 'r.', label=r'Approximated Posterior $2\\sigma$ Error Bars')\nplt.plot(true_beta, beta_loc - 2*jnp.exp(beta_log_scale), 'r.')\nplot_scale = 3\nplt.plot([-plot_scale, plot_scale], [-plot_scale, plot_scale], 'k')\nplt.xlabel('True beta')\nplt.ylabel('Estimated beta')\nplt.legend(loc='best')\n```\n\nExample:\n```text\n<matplotlib.legend.Legend at 0x73c6f434c080>\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.821Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":19,"totalLines":287,"estimatedTokens":1799}}109{"id":"doc-gpu_performance_tips_jax_documentation-8242fab7","source":"documentation","title":"GPU performance tips — JAX documentation","url":"https://docs.jax.dev/en/latest/gpu_performance_tips.html","text":"Example:\n```text\nimport os\nos.environ['XLA_FLAGS'] = (\n    '--xla_gpu_triton_gemm_any=True '\n    '--xla_gpu_enable_latency_hiding_scheduler=true '\n)\n```\n\nExample:\n```text\nexport JAX_ENABLE_PGLE=true\n\n# For JAX version <= 0.5.0 make sure to include:\nexport XLA_FLAGS=\"--xla_gpu_enable_latency_hiding_scheduler=true\"\n```\n\nExample:\n```text\nexport JAX_PGLE_PROFILING_RUNS=3\nexport JAX_PGLE_AGGREGATION_PERCENTILE=85\n\n# Right now the auto PGLE profile collection doesn't work with command buffer.\n# If the command buffer is enabled, Auto PGLE will disable it during profile\n# collection and enable it back after the recompilation. If you need to have a\n# consistent command buffer logic with and with PGLE profile you can disable it\n# manually:\nexport XLA_FLAGS=\"${XLA_FLAGS} --xla_gpu_enable_command_buffer=''\"\n```\n\nExample:\n```text\nimport jax\nfrom jax._src import config\n\nwith config.enable_pgle(True), config.pgle_profiling_runs(1):\n  # Run with the profiler collecting performance information.\n  train_step()\n  # Automatically re-compile with PGLE profile results\n  train_step()\n  ...\n```\n\nExample:\n```text\nimport jax\nfrom jax._src import config\n\ntrain_step_compiled = train_step().lower().compile()\n\nwith config.enable_pgle(True), config.pgle_profiling_runs(1):\n  train_step_compiled()\n  # No effect since module was pre-compiled.\n  train_step_compiled()\n```\n\nExample:\n```text\nexport JAX_ENABLE_COMPILATION_CACHE=yes          # not strictly needed, on by default\nexport JAX_COMPILATION_CACHE_DIR=/root/jax_cache\nJAX_ENABLE_PGLE=yes python my-model.py\n```\n\nExample:\n```text\nJAX_COMPILATION_CACHE_EXPECT_PGLE=yes nsys profile python my-model.py\n```\n\nExample:\n```text\nexport XLA_FLAGS=\"--xla_gpu_enable_latency_hiding_scheduler=true\"\n```\n\nExample:\n```text\nimport os\nfrom etils import epath\nimport jax\nfrom jax.experimental import profiler as exp_profiler\n\n# Define your profile directory\nprofile_dir = 'gs://my_bucket/profile'\njax.profiler.start_trace(profile_dir)\n\n# run your workflow\n# for i in range(10):\n#   train_step()\n\n# Stop trace\njax.profiler.stop_trace()\nprofile_dir = epath.Path(profile_dir)\ndirectories = profile_dir.glob('plugins/profile/*/')\ndirectories = [d for d in directories if d.is_dir()]\nrundir = directories[-1]\nlogging.info('rundir: %s', rundir)\n\n# Post process the profile\nfdo_profile = exp_profiler.get_profiled_instructions_proto(os.fspath(rundir))\n\n# Save the profile proto to a file.\ndump_dir = rundir / 'profile.pb'\ndump_dir.parent.mkdir(parents=True, exist_ok=True)\ndump_dir.write_bytes(fdo_profile)\n```\n\nExample:\n```text\nexport XLA_FLAGS=\"--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_pgle_profile_file_or_directory_path=/path/to/profile/profile.pb\"\n```\n\nExample:\n```text\nexport TF_CPP_MIN_LOG_LEVEL=0\n```\n\nExample:\n```text\n2023-07-21 16:09:43.551600: I external/xla/xla/service/gpu/gpu_hlo_schedule.cc:478] Using PGLE profile from /tmp/profile/plugins/profile/2023_07_20_18_29_30/profile.pb\n2023-07-21 16:09:43.551741: I external/xla/xla/service/gpu/gpu_hlo_schedule.cc:573] Found profile, using profile guided latency estimator\n```\n\nExample:\n```text\n--xla_gpu_enable_latency_hiding_scheduler=true\n--xla_gpu_enable_command_buffer=''\n--xla_disable_hlo_passes=collective-permute-motion\n--xla_gpu_experimental_pipeline_parallelism_opt_level=PIPELINE_PARALLELISM_OPT_LEVEL_ENABLE\n```\n\nExample:\n```text\n# Imports and setup\nimport functools\nimport jax\nfrom jax import sharding\nfrom jax.experimental import mesh_utils\nimport jax.numpy as jnp\nimport jax.random\n\nNUM_DEVICES = 4\nNUM_MICROBATCHES = 5\nNUM_CIRC_REPEATS = 2\nCONTRACTING_DIM_SIZE = 4096\nNON_CONTRACTING_DIM_SIZE = 8192\nCOMPUTE_INTENSITY = 32\n\n# Creates a collective permute for the \"forward edge\".\n# 0->1, 1->2, ... (N-2)->(N-1)\ndef shift_right(arr):\n  padding = [[1, 0]] + [[0, 0]] * (arr.ndim - 1)\n  # Use lax.slice to guarantee the gradient is a pad.\n  return jax.lax.slice(jnp.pad(arr, padding), [0] * arr.ndim, arr.shape)\n\n\n# Creates a collective permute for the \"back edge\".\n# (N-1)->0\ndef cycle_back(arr):\n  padding = [[0, NUM_DEVICES - 1]] + [[0, 0]] * (arr.ndim - 1)\n  return jax.lax.slice(\n      jnp.pad(arr, padding),\n      [NUM_DEVICES - 1] + [0] * (arr.ndim - 1),\n      (NUM_DEVICES - 1 + arr.shape[0],) + arr.shape[1:],\n  )\n\n\ndef select_on_first_device(then_value, else_value):\n  assert then_value.shape == else_value.shape\n  is_first_device = jax.lax.broadcasted_iota(\"int32\", then_value.shape, 0) == 0\n  return jnp.where(is_first_device, then_value, else_value)\n\n\ndef select_on_last_device(then_value, else_value):\n  assert then_value.shape == else_value.shape\n  is_last_device = (\n      jax.lax.broadcasted_iota(\"int32\", then_value.shape, 0) == NUM_DEVICES - 1\n  )\n  return jnp.where(is_last_device, then_value, else_value)\n\n\ndef select_on_first_cycle(i, then_value, else_value):\n  assert then_value.shape == else_value.shape\n  is_first_cycle = i < NUM_MICROBATCHES\n  return jnp.where(is_first_cycle, then_value, else_value)\n\n\ndef while_body(carry, i):\n  \"\"\"Body of the pipeline while loop.\"\"\"\n  weights, input_buffer, output_buffer, fwd_edge_data, bwd_edge_data = carry\n\n  # Read input data from input buffer.\n  input_data = jax.lax.dynamic_slice(\n      input_buffer,\n      (0, (i + 0) % NUM_MICROBATCHES, 0, 0),\n      (NUM_DEVICES, 1, CONTRACTING_DIM_SIZE, NON_CONTRACTING_DIM_SIZE),\n  )\n\n  # Collective permute on the \"forward edge\" shifts data to the next stage.\n  fwd_edge_data = shift_right(fwd_edge_data)\n\n  # Select compute argument based on device and pipeline cycle.\n  compute_argument = select_on_first_device(\n      select_on_first_cycle(i, input_data, bwd_edge_data),\n      fwd_edge_data,\n  ).reshape((NUM_DEVICES, CONTRACTING_DIM_SIZE, NON_CONTRACTING_DIM_SIZE))\n\n  # A few matmuls to simulate compute.\n  tmp = compute_argument\n  for _ in range(COMPUTE_INTENSITY):\n    tmp = jax.lax.dot_general(weights, tmp, (((2,), (1,)), ((0,), (0,))))\n  compute_result = tmp.reshape(\n      (NUM_DEVICES, 1, CONTRACTING_DIM_SIZE, NON_CONTRACTING_DIM_SIZE)\n  )\n\n  # Read data from buffer to pass it to the first device of the pipeline on the\n  # \"back edge\".\n  bwd_edge_data = jax.lax.dynamic_slice(\n      output_buffer,\n      (0, (1 + i) % NUM_MICROBATCHES, 0, 0),\n      (NUM_DEVICES, 1, CONTRACTING_DIM_SIZE, NON_CONTRACTING_DIM_SIZE),\n  )\n\n  # Collective permute on the \"back edge\" passes data to the first device.\n  bwd_edge_data = cycle_back(bwd_edge_data)\n\n  # Update output buffer. We do this after reading from it to avoid the data\n  # dependency.\n  output_buffer = jax.lax.dynamic_update_slice(\n      output_buffer,\n      compute_result,\n      (0, (2 + i) % NUM_MICROBATCHES, 0, 0),\n  )\n\n  fwd_edge_data = compute_result\n  carry = (\n      weights,\n      input_buffer,\n      output_buffer,\n      fwd_edge_data,\n      bwd_edge_data,\n  )\n  return carry, i\n\n\n@jax.jit(static_argnames=[\"mesh\"])\ndef entry_computation(weights, input_buffer, mesh):\n\n  # Init output buffer.\n  output_buffer = jnp.zeros_like(input_buffer)\n\n  # Init dummy data for forward and backward edge passed through the while loop.\n  dummy_data = jnp.zeros(\n      shape=(NUM_DEVICES, 1, CONTRACTING_DIM_SIZE, NON_CONTRACTING_DIM_SIZE)\n  ).astype(jnp.float32)\n  dummy_data = jax.device_put(\n      dummy_data,\n      sharding.NamedSharding(\n          mesh, sharding.PartitionSpec(\"x\")\n      ),\n  )\n\n  # Start pipeline.\n  carry = weights, input_buffer, output_buffer, dummy_data, dummy_data\n  num_iterations = NUM_CIRC_REPEATS * NUM_MICROBATCHES + NUM_DEVICES - 1\n  carry, _ = jax.lax.scan(while_body, carry, xs=jnp.arange(num_iterations))\n  _, _, output_buffer, _, _ = carry\n\n  return output_buffer\n\n\ndef main(_):\n\n  # Expect constant number of devices.\n  assert NUM_DEVICES == jax.local_device_count()\n\n  # Create mesh.\n  mesh = sharding.Mesh(\n      mesh_utils.create_device_mesh([NUM_DEVICES]),\n      axis_names=[\"x\"],\n  )\n\n  # Init weights.\n  weights = 1.0 / CONTRACTING_DIM_SIZE\n  weights = jax.lax.broadcast_in_dim(\n      weights,\n      shape=(NUM_DEVICES, CONTRACTING_DIM_SIZE, CONTRACTING_DIM_SIZE),\n      broadcast_dimensions=(),\n  )\n  weights = jax.device_put(\n      weights,\n      sharding.NamedSharding(\n          mesh, sharding.PartitionSpec(\"x\")\n      ),\n  )\n\n  # Init random input and replicate it across all devices.\n  random_key = jax.random.key(0)\n  input_buffer = jax.random.uniform(\n      random_key,\n      shape=(\n          NUM_MICROBATCHES,\n          CONTRACTING_DIM_SIZE,\n          NON_CONTRACTING_DIM_SIZE,\n      ),\n  )\n  input_buffer = jax.lax.broadcast_in_dim(\n      input_buffer,\n      shape=(\n          NUM_DEVICES,\n          NUM_MICROBATCHES,\n          CONTRACTING_DIM_SIZE,\n          NON_CONTRACTING_DIM_SIZE,\n      ),\n      broadcast_dimensions=[1, 2, 3],\n  )\n  input_buffer = jax.device_put(\n      input_buffer,\n      sharding.NamedSharding(\n          mesh, sharding.PartitionSpec(\"x\")\n      ),\n  )\n\n  # Run computation.\n  output_buffer = entry_computation(weights, input_buffer, mesh)\n  print(f\"output_buffer = \\n{output_buffer}\")\n```\n\nExample:\n```text\n## same setup and imports\ndef while_body(carry, i):\n  (\n      weights,\n      input_buffer,\n      output_buffer,\n      prev_compute_res,\n      prev_stage_slice_fwd,\n      prev_stage_slice_bwd,\n  ) = carry\n\n  # Read input data from input buffer.\n  input_slice = jax.lax.dynamic_slice(\n      input_buffer,\n      (0, (i + 0) % NUM_MICROBATCHES, 0, 0),\n      (1, 1, CONTRACTING_DIM_SIZE, NON_CONTRACTING_DIM_SIZE),\n  )\n\n  # send_fwd\n  fwd_send_token = jax.lax.psend(\n      prev_compute_res,\n      axis_name=\"x\",\n      perm=[(0, 1), (1, 2), (2, 3)],\n  )\n\n  # Select compute argument based on device and pipeline cycle\n  compute_argument = select_on_first_device(\n      select_on_first_cycle(i, input_slice, prev_stage_slice_bwd),\n      prev_stage_slice_fwd,\n  ).reshape((1, CONTRACTING_DIM_SIZE, NON_CONTRACTING_DIM_SIZE))\n\n  tmp = compute_argument\n  for _ in range(COMPUTE_INTENSITY):\n    tmp = jax.lax.dot_general(weights, tmp, (((2,), (1,)), ((0,), (0,))))\n  compute_result = tmp.reshape(\n      (1, 1, CONTRACTING_DIM_SIZE, NON_CONTRACTING_DIM_SIZE)\n  )\n\n  buffer_slice_for_bwd_ppermute = jax.lax.dynamic_slice(\n      output_buffer,\n      (0, (i + 1) % NUM_MICROBATCHES, 0, 0),\n      (1, 1, CONTRACTING_DIM_SIZE, NON_CONTRACTING_DIM_SIZE),\n  )\n\n  # make sure ppermute is scheduled after send_fwd\n  buffer_slice_for_bwd_ppermute_after_send_fwd, _ = (\n      jax.lax.optimization_barrier(\n          (buffer_slice_for_bwd_ppermute, fwd_send_token)\n      )\n  )\n  # ppermute_bwd\n  ppermute_bwd_data = jax.lax.ppermute(\n      buffer_slice_for_bwd_ppermute_after_send_fwd,\n      axis_name=\"x\",\n      perm=[(3, 0)],\n  )\n\n  # make sure recv is scheduled after ppermute\n  precv_token, _ = jax.lax.optimization_barrier(\n      (jax.lax.create_token(), ppermute_bwd_data)\n  )\n\n  # recv_fwd, matches the send_fwd in the next iteration\n  fwd_recv_data = jax.lax.precv(\n      precv_token,\n      out_shape=jax.ShapeDtypeStruct(\n          input_slice.shape, input_slice.dtype\n      ),\n      axis_name=\"x\",\n      perm=[(0, 1), (1, 2), (2, 3)],\n  )\n  update_output_buffer = jax.lax.dynamic_update_slice(\n      output_buffer,\n      compute_result,\n      (0, (i + 2) % NUM_MICROBATCHES, 0, 0),\n  )\n  carry = (\n      weights,\n      input_buffer,\n      update_output_buffer,\n      compute_result,\n      fwd_recv_data,\n      ppermute_bwd_data,\n  )\n  return carry, i\n\n\ndef entry_computation(\n    weights, input_buffer, dummy_data, mesh\n):\n\n  # Init output buffer.\n  output_buffer = jnp.zeros_like(input_buffer)\n\n  # Start pipeline.\n  dummy_slice_fwd = jax.lax.precv(\n      jax.lax.create_token(),\n      jax.ShapeDtypeStruct.like(dummy_data),\n      axis_name=\"x\",\n      perm=[(0, 1), (1, 2), (2, 3)],\n  )\n\n  carry = (\n      weights,\n      input_buffer,\n      output_buffer,\n      dummy_slice_fwd,\n      dummy_data,\n      dummy_data,\n  )\n\n  num_iterations = NUM_CIRC_REPEATS * NUM_MICROBATCHES + NUM_DEVICES - 1\n  carry, _ = jax.lax.scan(while_body, carry, xs=jnp.arange(num_iterations))\n\n  _ = jax.lax.psend(\n      carry[3],\n      axis_name=\"x\",\n      perm=[(0, 1), (1, 2), (2, 3)],\n  )\n\n  _, _, output_buffer, _, _, _ = carry\n\n  return output_buffer\n\n\ndef main(_):\n\n  # Expect constant number of devices.\n  assert NUM_DEVICES == jax.local_device_count()\n\n  # Create mesh.\n  mesh = Mesh(\n      mesh_utils.create_device_mesh([NUM_DEVICES]),\n      axis_names=[\"x\"],\n  )\n  # Init weights.\n  weights = 1.0 / CONTRACTING_DIM_SIZE\n  weights = jax.lax.broadcast_in_dim(\n      weights,\n      shape=(NUM_DEVICES, CONTRACTING_DIM_SIZE, CONTRACTING_DIM_SIZE),\n      broadcast_dimensions=(),\n  )\n  weights = jax.device_put(\n      weights, NamedSharding(mesh, P(\"x\"))\n  )\n  # Init input.\n  random_key = jax.random.key(0)\n  input_buffer = jax.random.uniform(\n      random_key,\n      shape=(\n          NUM_MICROBATCHES,\n          CONTRACTING_DIM_SIZE,\n          NON_CONTRACTING_DIM_SIZE,\n      ),\n  )\n  input_buffer = jax.lax.broadcast_in_dim(\n      input_buffer,\n      shape=(\n          NUM_DEVICES,\n          NUM_MICROBATCHES,\n          CONTRACTING_DIM_SIZE,\n          NON_CONTRACTING_DIM_SIZE,\n      ),\n      broadcast_dimensions=[1, 2, 3],\n  )\n\n  input_buffer = jax.device_put(\n      input_buffer,\n      NamedSharding(mesh, P(\"x\")),\n  )\n  # Init dummy data for forward and backward edge passed through the while\n  # loop.\n  dummy_slice = jnp.zeros(\n      shape=(NUM_DEVICES, 1, CONTRACTING_DIM_SIZE, NON_CONTRACTING_DIM_SIZE)\n  ).astype(jnp.float32)\n  dummy_data = jax.device_put(\n      dummy_slice,\n      NamedSharding(mesh, P(\"x\")),\n  )\n\n  entry = partial(entry_computation, mesh=mesh)\n\n  output_buffer = jax.jit(\n      jax.shard_map(\n          entry,\n          mesh=mesh,\n          in_specs=P(\"x\"),\n          out_specs=P(\"x\"),\n          check_vma=False,\n      )\n  )(weights, input_buffer, dummy_data)\n  print(f\"output_buffer = \\n{output_buffer}\")\n```\n\nExample:\n```text\nos.environ.update({\n  \"NCCL_LL128_BUFFSIZE\": \"-2\",\n  \"NCCL_LL_BUFFSIZE\": \"-2\",\n   \"NCCL_PROTO\": \"SIMPLE,LL,LL128\",\n })\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.822Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":16,"totalLines":532,"estimatedTokens":3483}}110{"id":"doc-training_a_simple_neural_network_with_tensorflow-ef2c7a4c","source":"documentation","title":"Training a simple neural network, with tensorflow/datasets data loading — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/neural_network_with_tfds_data.html","text":"Example:\n```text\nimport jax.numpy as jnp\nfrom jax import grad, jit, vmap\nfrom jax import random\n```\n\nExample:\n```text\n# A helper function to randomly initialize weights and biases\n# for a dense neural network layer\ndef random_layer_params(m, n, key, scale=1e-2):\n  w_key, b_key = random.split(key)\n  return scale * random.normal(w_key, (n, m)), scale * random.normal(b_key, (n,))\n\n# Initialize all layers for a fully-connected neural network with sizes \"sizes\"\ndef init_network_params(sizes, key):\n  keys = random.split(key, len(sizes))\n  return [random_layer_params(m, n, k) for m, n, k in zip(sizes[:-1], sizes[1:], keys)]\n\nlayer_sizes = [784, 512, 512, 10]\nstep_size = 0.01\nnum_epochs = 10\nbatch_size = 128\nn_targets = 10\nparams = init_network_params(layer_sizes, random.key(0))\n```\n\nExample:\n```text\nfrom jax.scipy.special import logsumexp\n\ndef relu(x):\n  return jnp.maximum(0, x)\n\ndef predict(params, image):\n  # per-example predictions\n  activations = image\n  for w, b in params[:-1]:\n    outputs = jnp.dot(w, activations) + b\n    activations = relu(outputs)\n\n  final_w, final_b = params[-1]\n  logits = jnp.dot(final_w, activations) + final_b\n  return logits - logsumexp(logits)\n```\n\nExample:\n```text\n# This works on single examples\nrandom_flattened_image = random.normal(random.key(1), (28 * 28,))\npreds = predict(params, random_flattened_image)\nprint(preds.shape)\n```\n\nExample:\n```text\n(10,)\n```\n\nExample:\n```text\n# Doesn't work with a batch\nrandom_flattened_images = random.normal(random.key(1), (10, 28 * 28))\ntry:\n  preds = predict(params, random_flattened_images)\nexcept TypeError:\n  print('Invalid shapes!')\n```\n\nExample:\n```text\nInvalid shapes!\n```\n\nExample:\n```text\n# Let's upgrade it to handle batches using `vmap`\n\n# Make a batched version of the `predict` function\nbatched_predict = vmap(predict, in_axes=(None, 0))\n\n# `batched_predict` has the same call signature as `predict`\nbatched_preds = batched_predict(params, random_flattened_images)\nprint(batched_preds.shape)\n```\n\nExample:\n```text\n(10, 10)\n```\n\nExample:\n```text\ndef one_hot(x, k, dtype=jnp.float32):\n  \"\"\"Create a one-hot encoding of x of size k.\"\"\"\n  return jnp.array(x[:, None] == jnp.arange(k), dtype)\n\ndef accuracy(params, images, targets):\n  target_class = jnp.argmax(targets, axis=1)\n  predicted_class = jnp.argmax(batched_predict(params, images), axis=1)\n  return jnp.mean(predicted_class == target_class)\n\ndef loss(params, images, targets):\n  preds = batched_predict(params, images)\n  return -jnp.mean(preds * targets)\n\n@jit\ndef update(params, x, y):\n  grads = grad(loss)(params, x, y)\n  return [(w - step_size * dw, b - step_size * db)\n          for (w, b), (dw, db) in zip(params, grads)]\n```\n\nExample:\n```text\nimport tensorflow as tf\n# Ensure TF does not see GPU and grab all GPU memory.\ntf.config.set_visible_devices([], device_type='GPU')\n\nimport tensorflow_datasets as tfds\n\ndata_dir = '/tmp/tfds'\n\n# Fetch full datasets for evaluation\n# tfds.load returns tf.Tensors (or tf.data.Datasets if batch_size != -1)\n# You can convert them to NumPy arrays (or iterables of NumPy arrays) with tfds.dataset_as_numpy\nmnist_data, info = tfds.load(name=\"mnist\", batch_size=-1, data_dir=data_dir, with_info=True)\nmnist_data = tfds.as_numpy(mnist_data)\ntrain_data, test_data = mnist_data['train'], mnist_data['test']\nnum_labels = info.features['label'].num_classes\nh, w, c = info.features['image'].shape\nnum_pixels = h * w * c\n\n# Full train set\ntrain_images, train_labels = train_data['image'], train_data['label']\ntrain_images = jnp.reshape(train_images, (len(train_images), num_pixels))\ntrain_labels = one_hot(train_labels, num_labels)\n\n# Full test set\ntest_images, test_labels = test_data['image'], test_data['label']\ntest_images = jnp.reshape(test_images, (len(test_images), num_pixels))\ntest_labels = one_hot(test_labels, num_labels)\n```\n\nExample:\n```text\nprint('Train:', train_images.shape, train_labels.shape)\nprint('Test:', test_images.shape, test_labels.shape)\n```\n\nExample:\n```text\nTrain: (60000, 784) (60000, 10)\nTest: (10000, 784) (10000, 10)\n```\n\nExample:\n```text\nimport time\n\ndef get_train_batches():\n  # as_supervised=True gives us the (image, label) as a tuple instead of a dict\n  ds = tfds.load(name='mnist', split='train', as_supervised=True, data_dir=data_dir)\n  # You can build up an arbitrary tf.data input pipeline\n  ds = ds.batch(batch_size).prefetch(1)\n  # tfds.dataset_as_numpy converts the tf.data.Dataset into an iterable of NumPy arrays\n  return tfds.as_numpy(ds)\n\nfor epoch in range(num_epochs):\n  start_time = time.time()\n  for x, y in get_train_batches():\n    x = jnp.reshape(x, (len(x), num_pixels))\n    y = one_hot(y, num_labels)\n    params = update(params, x, y)\n  epoch_time = time.time() - start_time\n\n  train_acc = accuracy(params, train_images, train_labels)\n  test_acc = accuracy(params, test_images, test_labels)\n  print(\"Epoch {} in {:0.2f} sec\".format(epoch, epoch_time))\n  print(\"Training set accuracy {}\".format(train_acc))\n  print(\"Test set accuracy {}\".format(test_acc))\n```\n\nExample:\n```text\nEpoch 0 in 28.30 sec\nTraining set accuracy 0.8400499820709229\nTest set accuracy 0.8469000458717346\nEpoch 1 in 14.74 sec\nTraining set accuracy 0.8743667006492615\nTest set accuracy 0.8803000450134277\nEpoch 2 in 14.57 sec\nTraining set accuracy 0.8901500105857849\nTest set accuracy 0.8957000374794006\nEpoch 3 in 14.36 sec\nTraining set accuracy 0.8991333246231079\nTest set accuracy 0.903700053691864\nEpoch 4 in 14.20 sec\nTraining set accuracy 0.9061833620071411\nTest set accuracy 0.9087000489234924\nEpoch 5 in 14.89 sec\nTraining set accuracy 0.9113333225250244\nTest set accuracy 0.912600040435791\nEpoch 6 in 13.95 sec\nTraining set accuracy 0.9156833291053772\nTest set accuracy 0.9176000356674194\nEpoch 7 in 13.32 sec\nTraining set accuracy 0.9192000031471252\nTest set accuracy 0.9214000701904297\nEpoch 8 in 13.55 sec\nTraining set accuracy 0.9222500324249268\nTest set accuracy 0.9241000413894653\nEpoch 9 in 13.40 sec\nTraining set accuracy 0.9253666996955872\nTest set accuracy 0.9269000291824341\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.824Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":15,"totalLines":219,"estimatedTokens":1506}}111{"id":"doc-training_a_simple_neural_network_with_pytorch_da-b76bb182","source":"documentation","title":"Training a simple neural network, with PyTorch data loading — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/Neural_Network_and_Data_Loading.html","text":"Example:\n```text\nimport jax.numpy as jnp\nfrom jax import grad, jit, vmap\nfrom jax import random\n```\n\nExample:\n```text\n# A helper function to randomly initialize weights and biases\n# for a dense neural network layer\ndef random_layer_params(m, n, key, scale=1e-2):\n  w_key, b_key = random.split(key)\n  return scale * random.normal(w_key, (n, m)), scale * random.normal(b_key, (n,))\n\n# Initialize all layers for a fully-connected neural network with sizes \"sizes\"\ndef init_network_params(sizes, key):\n  keys = random.split(key, len(sizes))\n  return [random_layer_params(m, n, k) for m, n, k in zip(sizes[:-1], sizes[1:], keys)]\n\nlayer_sizes = [784, 512, 512, 10]\nstep_size = 0.01\nnum_epochs = 8\nbatch_size = 128\nn_targets = 10\nparams = init_network_params(layer_sizes, random.key(0))\n```\n\nExample:\n```text\nfrom jax.scipy.special import logsumexp\n\ndef relu(x):\n  return jnp.maximum(0, x)\n\ndef predict(params, image):\n  # per-example predictions\n  activations = image\n  for w, b in params[:-1]:\n    outputs = jnp.dot(w, activations) + b\n    activations = relu(outputs)\n\n  final_w, final_b = params[-1]\n  logits = jnp.dot(final_w, activations) + final_b\n  return logits - logsumexp(logits)\n```\n\nExample:\n```text\n# This works on single examples\nrandom_flattened_image = random.normal(random.key(1), (28 * 28,))\npreds = predict(params, random_flattened_image)\nprint(preds.shape)\n```\n\nExample:\n```text\n(10,)\n```\n\nExample:\n```text\n# Doesn't work with a batch\nrandom_flattened_images = random.normal(random.key(1), (10, 28 * 28))\ntry:\n  preds = predict(params, random_flattened_images)\nexcept TypeError:\n  print('Invalid shapes!')\n```\n\nExample:\n```text\nInvalid shapes!\n```\n\nExample:\n```text\n# Let's upgrade it to handle batches using `vmap`\n\n# Make a batched version of the `predict` function\nbatched_predict = vmap(predict, in_axes=(None, 0))\n\n# `batched_predict` has the same call signature as `predict`\nbatched_preds = batched_predict(params, random_flattened_images)\nprint(batched_preds.shape)\n```\n\nExample:\n```text\n(10, 10)\n```\n\nExample:\n```text\ndef one_hot(x, k, dtype=jnp.float32):\n  \"\"\"Create a one-hot encoding of x of size k.\"\"\"\n  return jnp.array(x[:, None] == jnp.arange(k), dtype)\n\ndef accuracy(params, images, targets):\n  target_class = jnp.argmax(targets, axis=1)\n  predicted_class = jnp.argmax(batched_predict(params, images), axis=1)\n  return jnp.mean(predicted_class == target_class)\n\ndef loss(params, images, targets):\n  preds = batched_predict(params, images)\n  return -jnp.mean(preds * targets)\n\n@jit\ndef update(params, x, y):\n  grads = grad(loss)(params, x, y)\n  return [(w - step_size * dw, b - step_size * db)\n          for (w, b), (dw, db) in zip(params, grads)]\n```\n\nExample:\n```text\n!pip install torch torchvision\n```\n\nExample:\n```text\nRequirement already satisfied: torch in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (2.4.1)\nRequirement already satisfied: torchvision in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (0.19.1)\nRequirement already satisfied: filelock in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (from torch) (3.16.0)\nRequirement already satisfied: typing-extensions>=4.8.0 in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (from torch) (4.12.2)\nRequirement already satisfied: sympy in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (from torch) (1.13.2)\nRequirement already satisfied: networkx in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (from torch) (3.3)\nRequirement already satisfied: jinja2 in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (from torch) (3.1.4)\nRequirement already satisfied: fsspec in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (from torch) (2024.9.0)\nRequirement already satisfied: setuptools in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (from torch) (73.0.1)\nRequirement already satisfied: numpy in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (from torchvision) (1.26.4)\nRequirement already satisfied: pillow!=8.3.*,>=5.3.0 in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (from torchvision) (10.4.0)\nRequirement already satisfied: MarkupSafe>=2.0 in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (from jinja2->torch) (2.1.5)\nRequirement already satisfied: mpmath<1.4,>=1.1.0 in /home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages (from sympy->torch) (1.3.0)\n```\n\nExample:\n```text\n/home/m/.opt/miniforge3/envs/jax/lib/python3.12/pty.py:95: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.\n  pid, fd = os.forkpty()\n```\n\nExample:\n```text\nimport numpy as np\nfrom jax.tree_util import tree_map\nfrom torch.utils.data import DataLoader, default_collate\nfrom torchvision.datasets import MNIST\n\ndef numpy_collate(batch):\n  \"\"\"\n  Collate function specifies how to combine a list of data samples into a batch.\n  default_collate creates pytorch tensors, then tree_map converts them into numpy arrays.\n  \"\"\"\n  return tree_map(np.asarray, default_collate(batch))\n\ndef flatten_and_cast(pic):\n  \"\"\"Convert PIL image to flat (1-dimensional) numpy array.\"\"\"\n  return np.ravel(np.array(pic, dtype=jnp.float32))\n```\n\nExample:\n```text\n# Define our dataset, using torch datasets\nmnist_dataset = MNIST('/tmp/mnist/', download=True, transform=flatten_and_cast)\n# Create pytorch data loader with custom collate function\ntraining_generator = DataLoader(mnist_dataset, batch_size=batch_size, collate_fn=numpy_collate)\n```\n\nExample:\n```text\nDownloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\nFailed to download (trying next):\nHTTP Error 404: Not Found\n\nDownloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz\nDownloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to /tmp/mnist/MNIST/raw/train-images-idx3-ubyte.gz\nExtracting /tmp/mnist/MNIST/raw/train-images-idx3-ubyte.gz to /tmp/mnist/MNIST/raw\n\nDownloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\nFailed to download (trying next):\nHTTP Error 404: Not Found\n\nDownloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz\nDownloading https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz to /tmp/mnist/MNIST/raw/train-labels-idx1-ubyte.gz\nExtracting /tmp/mnist/MNIST/raw/train-labels-idx1-ubyte.gz to /tmp/mnist/MNIST/raw\n\nDownloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\nFailed to download (trying next):\nHTTP Error 404: Not Found\n\nDownloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz\nDownloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz to /tmp/mnist/MNIST/raw/t10k-images-idx3-ubyte.gz\nExtracting /tmp/mnist/MNIST/raw/t10k-images-idx3-ubyte.gz to /tmp/mnist/MNIST/raw\n\nDownloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\nFailed to download (trying next):\nHTTP Error 404: Not Found\n\nDownloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz\nDownloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz to /tmp/mnist/MNIST/raw/t10k-labels-idx1-ubyte.gz\nExtracting /tmp/mnist/MNIST/raw/t10k-labels-idx1-ubyte.gz to /tmp/mnist/MNIST/raw\n```\n\nExample:\n```text\n100.0%\n100.0%\n100.0%\n100.0%\n```\n\nExample:\n```text\n# Get the full train dataset (for checking accuracy while training)\ntrain_images = np.array(mnist_dataset.train_data).reshape(len(mnist_dataset.train_data), -1)\ntrain_labels = one_hot(np.array(mnist_dataset.train_labels), n_targets)\n\n# Get full test dataset\nmnist_dataset_test = MNIST('/tmp/mnist/', download=True, train=False)\ntest_images = jnp.array(mnist_dataset_test.test_data.numpy().reshape(len(mnist_dataset_test.test_data), -1), dtype=jnp.float32)\ntest_labels = one_hot(np.array(mnist_dataset_test.test_labels), n_targets)\n```\n\nExample:\n```text\n/home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages/torchvision/datasets/mnist.py:76: UserWarning: train_data has been renamed data\n  warnings.warn(\"train_data has been renamed data\")\n/home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages/torchvision/datasets/mnist.py:66: UserWarning: train_labels has been renamed targets\n  warnings.warn(\"train_labels has been renamed targets\")\n/home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages/torchvision/datasets/mnist.py:81: UserWarning: test_data has been renamed data\n  warnings.warn(\"test_data has been renamed data\")\n/home/m/.opt/miniforge3/envs/jax/lib/python3.12/site-packages/torchvision/datasets/mnist.py:71: UserWarning: test_labels has been renamed targets\n  warnings.warn(\"test_labels has been renamed targets\")\n```\n\nExample:\n```text\nimport time\n\nfor epoch in range(num_epochs):\n  start_time = time.time()\n  for x, y in training_generator:\n    y = one_hot(y, n_targets)\n    params = update(params, x, y)\n  epoch_time = time.time() - start_time\n\n  train_acc = accuracy(params, train_images, train_labels)\n  test_acc = accuracy(params, test_images, test_labels)\n  print(\"Epoch {} in {:0.2f} sec\".format(epoch, epoch_time))\n  print(\"Training set accuracy {}\".format(train_acc))\n  print(\"Test set accuracy {}\".format(test_acc))\n```\n\nExample:\n```text\nEpoch 0 in 5.53 sec\nTraining set accuracy 0.9156666994094849\nTest set accuracy 0.9199000000953674\nEpoch 1 in 1.13 sec\nTraining set accuracy 0.9370499849319458\nTest set accuracy 0.9383999705314636\nEpoch 2 in 1.12 sec\nTraining set accuracy 0.9490833282470703\nTest set accuracy 0.9467999935150146\nEpoch 3 in 1.21 sec\nTraining set accuracy 0.9568833708763123\nTest set accuracy 0.9532999992370605\nEpoch 4 in 1.17 sec\nTraining set accuracy 0.9631666541099548\nTest set accuracy 0.9574999809265137\nEpoch 5 in 1.17 sec\nTraining set accuracy 0.9675000309944153\nTest set accuracy 0.9615999460220337\nEpoch 6 in 1.11 sec\nTraining set accuracy 0.9709500074386597\nTest set accuracy 0.9652999639511108\nEpoch 7 in 1.17 sec\nTraining set accuracy 0.9736999869346619\nTest set accuracy 0.967199981212616\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.825Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":21,"totalLines":283,"estimatedTokens":2523}}112{"id":"doc-exporting_and_serializing_staged_out_computation-1ea0aac4","source":"documentation","title":"Exporting and serializing staged-out computations — JAX documentation","url":"https://docs.jax.dev/en/latest/export/export.html","text":"Example:\n```text\n>>> import re\n>>> import numpy as np\n>>> import jax\n>>> from jax import export\n\n>>> def f(x): return 2 * x * x\n\n\n>>> exported: export.Exported = export.export(jax.jit(f))(\n...    jax.ShapeDtypeStruct((), np.float32))\n\n>>> # You can inspect the Exported object\n>>> print(re.search(r\".*@main.*\", exported.mlir_module()).group(0))\n  func.func public @main(%arg0: tensor<f32> loc(\"x\")) -> (tensor<f32> {jax.result_info = \"result\"}) {\n\n>>> # And you can serialize the Exported to a bytearray.\n>>> serialized: bytearray = exported.serialize()\n\n>>> # The serialized function can later be rehydrated and called from\n>>> # another JAX computation, possibly in another process.\n>>> rehydrated_exp: export.Exported = export.deserialize(serialized)\n\n>>> def callee(y):\n...  return 3. * rehydrated_exp.call(y * 4.)\n\n>>> callee(1.)\nArray(96., dtype=float32)\n```\n\nExample:\n```text\n>>> import jax\n>>> from jax import export\n>>> from typing import Callable\n\n>>> def f(x): return 7 * x * x * x\n\n>>> # Serialize 3 levels of VJP along with the primal function\n>>> blob: bytearray = export.export(jax.jit(f))(1.).serialize(vjp_order=3)\n>>> rehydrated_f: Callable = export.deserialize(blob).call\n\n>>> rehydrated_f(0.1)  # 7 * 0.1^3\nArray(0.007, dtype=float32)\n\n>>> jax.grad(rehydrated_f)(0.1)  # 7*3 * 0.1^2\nArray(0.21000001, dtype=float32)\n\n>>> jax.grad(jax.grad(rehydrated_f))(0.1)  # 7*3*2 * 0.1\nArray(4.2, dtype=float32)\n\n>>> jax.grad(jax.grad(jax.grad(rehydrated_f)))(0.1)  # 7*3*2\nArray(42., dtype=float32)\n\n>>> jax.grad(jax.grad(jax.grad(jax.grad(rehydrated_f))))(0.1)  \nTraceback (most recent call last):\nValueError: No VJP is available\n```\n\nExample:\n```text\n>>> import jax\n>>> from jax import export\n>>> from jax import lax\n>>> from jax._src import core\n>>> from jax._src.interpreters import mlir\n>>> # Define a new primitive backed by a custom call\n>>> new_prim = core.Primitive(\"new_prim\")\n>>> _ = new_prim.def_abstract_eval(lambda x: x)\n>>> _ = mlir.register_lowering(new_prim, lambda ctx, o: mlir.custom_call(\"my_new_prim\", operands=[o], result_types=[o.type]).results)\n>>> print(jax.jit(new_prim.bind).lower(1.).compiler_ir())\nmodule @jit_bind attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {\n  func.func public @main(%arg0: tensor<f32>) -> (tensor<f32> {jax.result_info = \"result\"}) {\n    %0 = stablehlo.custom_call @my_new_prim(%arg0) {api_version = 2 : i32, backend_config = \"\"} : (tensor<f32>) -> tensor<f32>\n    return %0 : tensor<f32>\n  }\n}\n\n>>> # If we try to export, we get an error\n>>> export.export(jax.jit(new_prim.bind))(1.)  \nTraceback (most recent call last):\nValueError: Cannot serialize code with custom calls whose targets have no compatibility guarantees: my_new_bind\n\n>>> # We can avoid the error if we pass a `DisabledSafetyCheck.custom_call`\n>>> exp = export.export(\n...    jax.jit(new_prim.bind),\n...    disabled_checks=[export.DisabledSafetyCheck.custom_call(\"my_new_prim\")])(1.)\n```\n\nExample:\n```text\n>>> from jax import export\n>>> export.default_export_platform()\n'cpu'\n```\n\nExample:\n```text\n>>> import jax\n>>> from jax import export\n>>> from jax import lax\n\n>>> # You can specify the export platform, e.g., `tpu`, `cpu`, `cuda`, `rocm`\n>>> # even if the current machine does not have that accelerator.\n>>> exp = export.export(jax.jit(lax.cos), platforms=['tpu'])(1.)\n\n>>> # But you will get an error if you try to compile `exp`\n>>> # on a machine that does not have TPUs.\n>>> exp.call(1.)  \nTraceback (most recent call last):\nValueError: Function 'cos' was lowered for platforms '('tpu',)' but it is used on '('cpu',)'.\n\n>>> # We can avoid the error if we pass a `DisabledSafetyCheck.platform`\n>>> # parameter to `export`, e.g., because you have reasons to believe\n>>> # that the code lowered will run adequately on the current\n>>> # compilation platform (which is the case for `cos` in this\n>>> # example):\n>>> exp_unsafe = export.export(jax.jit(lax.cos),\n...    platforms=['tpu'],\n...    disabled_checks=[export.DisabledSafetyCheck.platform()])(1.)\n\n>>> exp_unsafe.call(1.)\nArray(0.5403023, dtype=float32, weak_type=True)\n\n# and similarly with multi-platform lowering\n>>> exp_multi = export.export(jax.jit(lax.cos),\n...    platforms=['tpu', 'cpu', 'cuda'])(1.)\n>>> exp_multi.call(1.)\nArray(0.5403023, dtype=float32, weak_type=True)\n```\n\nExample:\n```text\n>>> import jax\n>>> from jax import export\n>>> from jax import lax\n>>> # A largish function\n>>> def f(x):\n...   for i in range(1000):\n...     x = jnp.cos(x)\n...   return x\n\n>>> exp_single = export.export(jax.jit(f))(1.)\n>>> len(exp_single.mlir_module_serialized)  \n9220\n\n>>> exp_multi = export.export(jax.jit(f),\n...                           platforms=[\"cpu\", \"tpu\", \"cuda\"])(1.)\n>>> len(exp_multi.mlir_module_serialized)  \n9282\n```\n\nExample:\n```text\n>>> import jax\n>>> from jax import export\n>>> from jax.sharding import AbstractMesh, Mesh, NamedSharding\n>>> from jax.sharding import PartitionSpec as P\n>>>\n>>> # Use an AbstractMesh for exporting\n>>> export_mesh = AbstractMesh((4,), (\"a\",))\n\n>>> def f(x):\n...   return x.T\n\n>>> exp = export.export(jax.jit(f))(\n...    jax.ShapeDtypeStruct((32,), dtype=np.int32,\n...                         sharding=NamedSharding(export_mesh, P(\"a\"))))\n\n>>> # `exp` knows for how many devices it was exported.\n>>> exp.nr_devices\n4\n\n>>> # and it knows the shardings for the inputs. These will be applied\n>>> # when the exported is called.\n>>> exp.in_shardings_jax(export_mesh)\n(NamedSharding(mesh=AbstractMesh('a': 4, axis_types=(Explicit,)), spec=P('a',)),)\n\n>>> # You can also use a concrete set of devices for exporting\n>>> concrete_devices = jax.local_devices()[:4]\n>>> concrete_mesh = Mesh(concrete_devices, (\"a\",))\n>>> exp2 = export.export(jax.jit(f))(\n...    jax.ShapeDtypeStruct((32,), dtype=np.int32,\n...                         sharding=NamedSharding(concrete_mesh, P(\"a\"))))\n\n>>> # When you call an Exported, you must use a concrete set of devices\n>>> arg = jax.device_put(jnp.arange(8 * 4),\n...                      NamedSharding(concrete_mesh, P(\"a\")))\n\n>>> res1 = exp.call(arg)\n>>> # Check out the first 2 shards of the result\n>>> [f\"device={s.device} index={s.index}\" for s in res1.addressable_shards[:2]]\n['device=cpu:0 index=(slice(0, 8, None),)', 'device=cpu:1 index=(slice(8, 16, None),)']\n\n>>> # We can call `exp` with some other 4 devices and another\n>>> # mesh with a different shape, as long as the number of devices is\n>>> # the same.\n>>> other_mesh = Mesh(np.array(jax.local_devices()[2:6]).reshape((2, 2)), (\"b\", \"c\"))\n>>> res2 = exp.call(jax.device_put(arg,\n...                                NamedSharding(other_mesh, P(\"b\"))))\n\n>>> # Check out the first 2 shards of the result. Notice that the output is\n>>> # sharded similarly; this means that the input was resharded according to the\n>>> # exp.in_shardings.\n>>> [f\"device={s.device} index={s.index}\" for s in res2.addressable_shards[:2]]\n['device=cpu:2 index=(slice(0, 8, None),)', 'device=cpu:3 index=(slice(8, 16, None),)']\n```\n\nExample:\n```text\n>>> import jax\n>>> from jax import export\n>>> from jax.sharding import Mesh, NamedSharding\n>>> from jax.sharding import PartitionSpec as P\n\n>>> export_devices = jax.local_devices()\n>>> export_mesh = Mesh(np.array(export_devices), (\"a\",))\n>>> def f(x):\n...   return x.T\n\n>>> exp = export.export(jax.jit(f))(\n...    jax.ShapeDtypeStruct((4 * len(export_devices),), dtype=np.int32,\n...                         sharding=NamedSharding(export_mesh, P(\"a\"))))\n\n>>> arg = jnp.arange(4 * len(export_devices))\n>>> exp.call(arg)  \nTraceback (most recent call last):\nValueError: Exported module f was lowered for 8 devices and is called in a context with 1 devices. This is disallowed because: the module was lowered for more than 1 device.\n```\n\nExample:\n```text\n>>> import jax\n>>> from jax import export\n>>> from jax.sharding import Mesh, NamedSharding\n>>> from jax.sharding import PartitionSpec as P\n\n>>> export_devices = jax.local_devices()\n>>> export_mesh = Mesh(np.array(export_devices), (\"a\",))\n>>> def f(x):\n...   return x.T\n\n\n>>> exp = export.export(jax.jit(f))(\n...    jax.ShapeDtypeStruct((4 * len(export_devices),), dtype=np.int32,\n...                         sharding=NamedSharding(export_mesh, P(\"a\"))))\n\n>>> # Prepare the mesh for calling `exp`.\n>>> calling_mesh = Mesh(np.array(export_devices[::-1]), (\"a\",))\n\n>>> # Shard the arg according to what `exp` expects.\n>>> arg = jnp.arange(4 * len(export_devices))\n>>> sharded_arg = jax.device_put(arg, exp.in_shardings_jax(calling_mesh)[0])\n>>> res = exp.call(sharded_arg)\n```\n\nExample:\n```text\n>>> import jax\n>>> from jax import export\n>>> from jax.sharding import Mesh, NamedSharding\n>>> from jax.sharding import PartitionSpec as P\n\n>>> def f(x):\n...   return jnp.cos(x)\n\n>>> arg = jnp.arange(4)\n>>> exp = export.export(jax.jit(f))(arg)\n>>> exp.in_avals\n(ShapedArray(int32[4]),)\n\n>>> exp.nr_devices\n1\n\n>>> # Prepare the mesh for calling `exp`.\n>>> calling_mesh = Mesh(jax.local_devices()[:4], (\"b\",))\n\n>>> # Shard the arg according to what `exp` expects.\n>>> sharded_arg = jax.device_put(arg,\n...                              NamedSharding(calling_mesh, P(\"b\")))\n>>> res = exp.call(sharded_arg)\n```\n\nExample:\n```text\n>>> from jax import export\n>>> exp: export.Exported = export.export(jnp.cos)(1.)\n>>> exp.calling_convention_version\n10\n```\n\nExample:\n```text\n>>> from jax import export\n>>> (export.minimum_supported_calling_convention_version, export.maximum_supported_calling_convention_version)\n(9, 10)\n\n>>> from jax._src import config\n>>> with config.jax_export_calling_convention_version(10):\n...  exp = export.export(jnp.cos)(1.)\n...  exp.calling_convention_version\n10\n```\n\nExample:\n```text\nfunc public main(\n            platform_index: i32 {jax.global_constant=\"_platform_index\"},\n            token_in: token,\n            arg: f32[?, ?]) {\n         arg_w = hlo.get_dimension_size(arg, 0)\n         dim1 = hlo.get_dimension_size(arg, 1)\n         arg_h = hlo.floordiv(dim1, 2)\n         call _check_shape_assertions(arg)  # See below\n         token = new_token()\n         token_out, res = call _wrapped_jax_export_main(platform_index,\n                                                        arg_h,\n                                                        arg_w,\n                                                        token_in,\n                                                        arg)\n         return token_out, res\n      }\n```\n\nExample:\n```text\nfunc private _wrapped_jax_export_main(\n          platform_index: i32 {jax.global_constant=\"_platform_index\"},\n          arg_h: i32 {jax.global_constant=\"h\"},\n          arg_w: i32 {jax.global_constant=\"w\"},\n          arg_token: stablehlo.token {jax.token=True},\n          arg: f32[?, ?]) -> (stablehlo.token, ...)\n```\n\nExample:\n```text\nfunc private _check_shape_assertions(arg: f32[?, ?]) {\n         # Check that w is >= 1\n         arg_w = hlo.get_dimension_size(arg, 0)\n         custom_call @shape_assertion(arg_w >= 1, arg_w,\n            error_message=\"Dimension variable 'w' must have integer value >= 1. Found {0}\")\n         # Check that dim1 is even\n         dim1 = hlo.get_dimension_size(arg, 1)\n         custom_call @shape_assertion(dim1 % 2 == 0, dim1 % 2,\n            error_message=\"Division had remainder {0} when computing the value of 'h')\n         # Check that h >= 1\n         arg_h = hlo.floordiv(dim1, 2)\n         custom_call @shape_assertion(arg_h >= 1, arg_h,\n            error_message=\"\"Dimension variable 'h' must have integer value >= 1. Found {0}\")\n```\n\nExample:\n```text\n# Log from python\npython tests/export_test.py JaxExportTest.test_basic -v=3\n# Or, log from pytest to /tmp/mylog.txt\npytest tests/export_test.py -k test_basic --log-level=3 --log-file=/tmp/mylog.txt\n```\n\nExample:\n```text\nI0619 10:54:18.978733 8299482112 _export.py:606] Exported JAX function: fun_name=sin version=9 lowering_platforms=('cpu',) disabled_checks=()\nI0619 10:54:18.978767 8299482112 _export.py:607] Define JAX_DUMP_IR_TO to dump the module.\n```\n\nExample:\n```text\nJAX_DUMP_IR_TO=/tmp/export.dumps pytest tests/export_test.py -k test_basic --log-level=3 --log-file=/tmp/mylog.txt\nINFO     absl:_export.py:606 Exported JAX function: fun_name=sin version=9 lowering_platforms=('cpu',) disabled_checks=()\nINFO     absl:_export.py:607 The module was dumped to jax_ir0_jit_sin_export.mlir.\n```\n\nExample:\n```text\n$ ls -l /tmp/export.dumps/\ntotal 32\n-rw-rw-r--@ 1 necula  wheel  2316 Jun 19 11:04 jax_ir0_jit_sin_export.mlir\n-rw-rw-r--@ 1 necula  wheel  2279 Jun 19 11:04 jax_ir1_jit_sin_compile.mlir\n-rw-rw-r--@ 1 necula  wheel  3377 Jun 19 11:04 jax_ir2_jit_call_exported_compile.mlir\n-rw-rw-r--@ 1 necula  wheel  2333 Jun 19 11:04 jax_ir3_jit_my_fun_export.mlir\n```\n\nExample:\n```text\nfrom jax._src import config\nfrom jax._src.lib import version as jaxlib_version\n\ndef my_lowering_rule(ctx: LoweringRuleContext, ...):\n  if ctx.is_forward_compat() or jaxlib_version < (0, 4, 31):\n    # this is the old lowering, using target T, while we\n    # are in forward compatibility mode for T, or we\n    # are in OSS and are using an old jaxlib.\n    return hlo.custom_call(\"T\", ...)\n  else:\n    # This is the new lowering, using target T_NEW, for\n    # when we use a jaxlib with version `>= (0, 4, 31)`\n    # (or when this is internal usage), and also we are\n    # in JIT mode.\n    return hlo.custom_call(\"T_NEW\", ...)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.827Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":20,"totalLines":403,"estimatedTokens":3325}}113{"id":"doc-fault_tolerant_distributed_jax_jax_documentation-fbd31113","source":"documentation","title":"Fault Tolerant Distributed JAX — JAX documentation","url":"https://docs.jax.dev/en/latest/fault_tolerance.html","text":"Example:\n```text\n1from absl import app\n 2from absl import flags\n 3from collections.abc import Sequence\n 4import jax\n 5import time\n 6\n 7_PROCESS_ID = flags.DEFINE_integer(\"i\", -1, \"Process id\")\n 8_NUM_PROCESSES = flags.DEFINE_integer(\"n\", -1, \"Number of processes\")\n 9\n10\n11def main(_: Sequence[str]) -> None:\n12  jax.distributed.initialize(\n13      coordinator_address=\"localhost:9000\",\n14      num_processes=_NUM_PROCESSES.value,\n15      process_id=_PROCESS_ID.value,\n16      local_device_ids=[_PROCESS_ID.value],\n17      heartbeat_timeout_seconds=10,\n18  )\n19  print(f'{jax.devices()=}')\n20  print(f'{jax.local_devices()=}')\n21  while True:\n22    print(time.time())\n23    time.sleep(1)\n24\n25\n26if __name__ == \"__main__\":\n27  app.run(main)\n```\n\nExample:\n```text\npython example.py --i=0 --n=4  # in terminal 1\npython example.py --i=1 --n=4  # in terminal 2\npython example.py --i=2 --n=4  # in terminal 3\npython example.py --i=3 --n=4  # in terminal 4\n```\n\nExample:\n```text\nE0926 17:26:32.075402  157988 coordination_service_agent.cc:332] Polled an error from coordination service (this can be an error from this or another task).\nF0926 17:26:32.075587  157988 client.h:77] Terminating process because the JAX distributed service detected fatal errors. This most likely indicates that another task died; see the other task logs for more details. Disable Python buffering, i.e. `python -u`, to be sure to see all the previous output. absl::Status: UNAVAILABLE: The following tasks are unhealthy (stopped sending heartbeats):\n/job:jax_worker/replica:0/task:3\nThe tasks have crashed. Check the task logs for an earlier error, or scheduler events (e.g. preemption, eviction) to debug further.\n\nRPC: /tensorflow.CoordinationService/PollForError [type.googleapis.com/tensorflow.CoordinationServiceError='']\n```\n\nExample:\n```text\n1import os\n 2os.environ['XLA_FLAGS'] = '--xla_gpu_nccl_terminate_on_error=false'\n 3\n 4from absl import app\n 5from absl import flags\n 6from collections.abc import Sequence\n 7import jax\n 8import time\n 9\n10_PROCESS_ID = flags.DEFINE_integer(\"i\", -1, \"Process id\")\n11_NUM_PROCESSES = flags.DEFINE_integer(\"n\", -1, \"Number of processes\")\n12\n13\n14def main(_: Sequence[str]) -> None:\n15  jax.config.update(\"jax_enable_recoverability\", True)\n16  jax.distributed.initialize(\n17      coordinator_address=\"localhost:9000\",\n18      num_processes=_NUM_PROCESSES.value,\n19      process_id=_PROCESS_ID.value,\n20      local_device_ids=[_PROCESS_ID.value],\n21      heartbeat_timeout_seconds=10,\n22  )\n23  print(f'{jax.devices()=}')\n24  print(f'{jax.local_devices()=}')\n25  while True:\n26    print(time.time())\n27    time.sleep(1)\n28\n29\n30if __name__ == \"__main__\":\n31  app.run(main)\n```\n\nExample:\n```text\nE0929 17:42:48.594192 1044529 coordination_service_agent.cc:332] Polled an error from coordination service (this can be an error from this or another task).\nF0929 17:42:48.594200 1044529 client.h:77] Terminating process because the JAX distributed service detected fatal errors. This most likely indicates that another task died; see the other task logs for more details. Disable Python buffering, i.e. `python -u`, to be sure to see all the previous output. absl::Status: UNAVAILABLE: Failed to send RPC to coordination service. Either the leader task was preempted/died/restarted unexpectedly or this task is experiencing network issues. Check earlier logs from 1) this task, 2) the leader (usually slice 0 task 0), and 3) cluster scheduler to debug further.\nAdditional GRPC error information from remote target coordination_service while calling /tensorflow.CoordinationService/PollForError:\n:UNKNOWN:Error received from peer  {grpc_message:\"Socket closed\", grpc_status:14}\n```\n\nExample:\n```text\n1import os\n 2os.environ['XLA_FLAGS'] = '--xla_gpu_nccl_terminate_on_error=false'\n 3\n 4from absl import app\n 5from absl import flags\n 6from collections.abc import Sequence\n 7import jax\n 8import jax.numpy as jnp\n 9import time\n10\n11_PROCESS_ID = flags.DEFINE_integer(\"i\", -1, \"Process id\")\n12_NUM_PROCESSES = flags.DEFINE_integer(\"n\", -1, \"Number of processes\")\n13\n14\n15def main(_: Sequence[str]) -> None:\n16  jax.config.update(\"jax_enable_recoverability\", True)\n17  jax.distributed.initialize(\n18      coordinator_address=\"localhost:9000\",\n19      num_processes=_NUM_PROCESSES.value,\n20      process_id=_PROCESS_ID.value,\n21      local_device_ids=[_PROCESS_ID.value],\n22      heartbeat_timeout_seconds=10,\n23  )\n24  print(f'{jax.devices()=}')\n25  print(f'{jax.local_devices()=}')\n26\n27  n = jax.device_count()\n28  jax.set_mesh(jax.make_mesh((n,), (\"i\",)))\n29  x = jax.device_put(jnp.arange(n), jax.P(\"i\"))\n30  while True:\n31    print(jnp.sum(x))\n32    time.sleep(1)\n33\n34\n35if __name__ == \"__main__\":\n36  app.run(main)\n```\n\nExample:\n```text\n1import os\n 2os.environ['XLA_FLAGS'] = ' '.join([\n 3  '--xla_gpu_nccl_terminate_on_error=false',\n 4  '--xla_gpu_nccl_async_execution=true',\n 5  '--xla_gpu_nccl_blocking_communicators=false',\n 6])\n 7os.environ['XLA_PYTHON_CLIENT_ABORT_COLLECTIVES_ON_FAILURE'] = '1'\n 8os.environ['XLA_PYTHON_CLIENT_USE_TFRT_GPU_CLIENT'] = '1'\n 9\n10from absl import app\n11from absl import flags\n12from collections.abc import Sequence\n13import jax\n14import jax.numpy as jnp\n15import time\n16\n17_PROCESS_ID = flags.DEFINE_integer(\"i\", -1, \"Process id\")\n18_NUM_PROCESSES = flags.DEFINE_integer(\"n\", -1, \"Number of processes\")\n19\n20\n21def main(_: Sequence[str]) -> None:\n22  jax.config.update(\"jax_enable_recoverability\", True)\n23  jax.distributed.initialize(\n24      coordinator_address=\"localhost:9000\",\n25      num_processes=_NUM_PROCESSES.value,\n26      process_id=_PROCESS_ID.value,\n27      local_device_ids=[_PROCESS_ID.value],\n28      heartbeat_timeout_seconds=10,\n29  )\n30  print(f'{jax.devices()=}')\n31  print(f'{jax.local_devices()=}')\n32\n33  # Don't do this. Use live_devices instead.\n34  from jax.experimental.multihost_utils import _live_devices\n35  _live_devices(jax._src.distributed.global_state.client, jax.devices())\n36\n37  n = jax.device_count()\n38  jax.set_mesh(jax.make_mesh((n,), (\"i\",)))\n39  x = jax.device_put(jnp.arange(n), jax.P(\"i\"))\n40  while True:\n41    print(jnp.sum(x))\n42    time.sleep(1)\n43\n44\n45if __name__ == \"__main__\":\n46  app.run(main)\n```\n\nExample:\n```text\njaxlib._jax.XlaRuntimeError: FAILED_PRECONDITION: Task with incarnation id 3446767950926952685 is not connected\n```\n\nExample:\n```text\n1import os\n 2os.environ['XLA_FLAGS'] = ' '.join([\n 3  '--xla_gpu_nccl_terminate_on_error=false',\n 4  '--xla_gpu_nccl_async_execution=true',\n 5  '--xla_gpu_nccl_blocking_communicators=false',\n 6])\n 7os.environ['XLA_PYTHON_CLIENT_ABORT_COLLECTIVES_ON_FAILURE'] = '1'\n 8os.environ['XLA_PYTHON_CLIENT_USE_TFRT_GPU_CLIENT'] = '1'\n 9\n10from absl import app\n11from absl import flags\n12from collections.abc import Sequence\n13from jax.experimental.multihost_utils import live_devices\n14import jax\n15import jax.numpy as jnp\n16import time\n17\n18_PROCESS_ID = flags.DEFINE_integer(\"i\", -1, \"Process id\")\n19_NUM_PROCESSES = flags.DEFINE_integer(\"n\", -1, \"Number of processes\")\n20\n21\n22def main(_: Sequence[str]) -> None:\n23  jax.config.update(\"jax_enable_recoverability\", True)\n24  jax.distributed.initialize(\n25      coordinator_address=\"localhost:9000\",\n26      num_processes=_NUM_PROCESSES.value,\n27      process_id=_PROCESS_ID.value,\n28      local_device_ids=[_PROCESS_ID.value],\n29      heartbeat_timeout_seconds=10,\n30  )\n31  print(f'{jax.devices()=}')\n32  print(f'{jax.local_devices()=}')\n33\n34  while True:\n35    try:\n36      with live_devices(jax.devices()) as devices:\n37        print(f'{devices=}')\n38        n = len(devices)\n39        jax.set_mesh(jax.make_mesh((n,), (\"i\",), devices=devices))\n40        x = jax.device_put(jnp.arange(n), jax.P(\"i\"))\n41        print(jnp.sum(x))\n42    except Exception as e:\n43      print('FAIL:', e)\n44    else:\n45      print('PASS')\n46    time.sleep(1)\n47\n48\n49if __name__ == \"__main__\":\n50  app.run(main)\n```\n\nExample:\n```text\ntry:\n  with live_devices(jax.devices()) as devices:\n    ...\nexcept Exception as e:\n  ... # Branch A\nelse:\n  ... # Branch B\n```\n\nExample:\n```text\nx = ...\ny = ...\ntry:\n  with live_devices(jax.devices()) as devices:\n    y = jnp.sum(x)\nexcept Exception as e:\n  ... # Branch A\nelse:\n  ... # Branch B\nprint(y)\n```\n\nExample:\n```text\nx = ...\ny = ...\ntry:\n  with live_devices(jax.devices()) as devices:\n    y = jax.block_until_ready(jnp.sum(x))\nexcept Exception as e:\n  ... # Branch A\nelse:\n  ... # Branch B\nprint(y)\n```\n\nExample:\n```text\n1import os\n 2os.environ['XLA_FLAGS'] = ' '.join([\n 3    '--xla_gpu_nccl_terminate_on_error=false',\n 4    '--xla_gpu_nccl_async_execution=true',\n 5    '--xla_gpu_nccl_blocking_communicators=false',\n 6])\n 7os.environ['XLA_PYTHON_CLIENT_ABORT_COLLECTIVES_ON_FAILURE'] = '1'\n 8os.environ['XLA_PYTHON_CLIENT_USE_TFRT_GPU_CLIENT'] = '1'\n 9\n10from absl import app\n11from absl import flags\n12from collections.abc import Sequence\n13from jax.experimental.multihost_utils import live_devices\n14import jax\n15import jax.numpy as jnp\n16import time\n17\n18_PROCESS_ID = flags.DEFINE_integer(\"i\", -1, \"Process id\")\n19_NUM_PROCESSES = flags.DEFINE_integer(\"n\", -1, \"Number of processes\")\n```\n\nExample:\n```text\n21def replicated(x: jax.Array, devices: list[jax.Device]):\n22  \"\"\"Return x replicated across the provided devices.\n23\n24  Note that replicated(x) doesn't actually move any data. It simply creates a\n25  logically replicated array with x as the local replica.\n26  \"\"\"\n27  n = len(devices)\n28  mesh = jax.make_mesh((n, ), (\"i\", ), devices=devices)\n29  spec = jax.sharding.PartitionSpec(None)\n30  sharding = jax.sharding.NamedSharding(mesh, spec)\n31  shards = [\n32      jax.device_put(x.addressable_shards[0].data, d) for d in devices\n33      if d.process_index == jax.process_index()\n34  ]\n35  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)\n```\n\nExample:\n```text\n38def sharded(x: jax.Array, devices: list[jax.Device]):\n39  \"\"\"Return x sharded across the provided devices.\n40\n41  Note that sharded(x) doesn't actually move any data. It simply creates a\n42  logically sharded array. x should have the same shape as the global array.\n43  \"\"\"\n44  n = len(devices)\n45  mesh = jax.make_mesh((n, ), (\"i\", ), devices=devices)\n46  spec = jax.sharding.PartitionSpec(\"i\")\n47  sharding = jax.sharding.NamedSharding(mesh, spec)\n48  m = sharding.addressable_devices_indices_map(x.shape)\n49  shards = [jax.device_put(x[m[d]], d) for d in jax.local_devices()]\n50  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)\n```\n\nExample:\n```text\n53def main(_: Sequence[str]) -> None:\n54  # Parse command line arguments and initialize multi-controller JAX.\n55  jax.config.update(\"jax_enable_recoverability\", True)\n56  jax.distributed.initialize(coordinator_address=\"localhost:8000\",\n57                             process_id=_PROCESS_ID.value,\n58                             num_processes=_NUM_PROCESSES.value,\n59                             local_device_ids=[_PROCESS_ID.value],\n60                             heartbeat_timeout_seconds=10)\n61  print(f'{jax.devices()=}')\n62  print(f'{jax.local_devices()=}')\n```\n\nExample:\n```text\n64  # Initialize the model's weights.\n65  keys = iter(jax.random.split(jax.random.key(seed=42), num=3))\n66  weights = jax.random.normal(next(keys), shape=(1, ))\n67\n68  # We'll learn a trivial linear model: a*x.\n69  def predict(weights, X):\n70    return weights * X\n71\n72  # We'll use mean squared error loss.\n73  def loss(weights, X, Y):\n74    return jnp.mean((predict(weights, X) - Y)**2)\n75\n76  # Initialize the (noisy) training data with a=10.\n77  X = jax.random.permutation(next(keys), jnp.arange(-300., 300.))\n78  Y = 10 * X + jax.random.normal(next(keys), X.shape)\n79\n80  # Hyperparameters.\n81  loss_and_grad = jax.jit(jax.value_and_grad(loss))\n82  learning_rate = 1e-6\n83  device_batch_size = 10\n```\n\nExample:\n```text\n85  step = 0\n 86  while True:\n 87    try:\n 88      with live_devices(jax.devices()) as devices:\n 89        print(f'=== Running step {step} with live devices = {devices} ===')\n 90\n 91        # Replicate the model weights.\n 92        weights = replicated(weights, devices)\n 93\n 94        # Shard the batch.\n 95        batch_size = device_batch_size * len(devices)\n 96        start = (step * batch_size) % len(X)\n 97        stop = start + batch_size\n 98        X_batch = sharded(X[start:stop], devices)\n 99        Y_batch = sharded(Y[start:stop], devices)\n100\n101        # Compute gradients and update weights.\n102        l, grad = loss_and_grad(weights, X_batch, Y_batch)\n103        new_weights = jax.block_until_ready(weights - learning_rate * grad)\n104    except Exception as e:\n105      print(f'Step {step} failed: {e}')\n106    else:\n107      print(f'Step {step} succeeded: loss = {l}')\n108      step += 1\n109      weights = new_weights\n110\n111    time.sleep(1)\n```\n\nExample:\n```text\n1import os\n  2os.environ['XLA_FLAGS'] = ' '.join([\n  3    '--xla_gpu_nccl_terminate_on_error=false',\n  4    '--xla_gpu_nccl_async_execution=true',\n  5    '--xla_gpu_nccl_blocking_communicators=false',\n  6])\n  7os.environ['XLA_PYTHON_CLIENT_ABORT_COLLECTIVES_ON_FAILURE'] = '1'\n  8os.environ['XLA_PYTHON_CLIENT_USE_TFRT_GPU_CLIENT'] = '1'\n  9\n 10from absl import app\n 11from absl import flags\n 12from collections.abc import Sequence\n 13from jax.experimental.multihost_utils import live_devices\n 14import jax\n 15import jax.numpy as jnp\n 16import time\n 17\n 18_PROCESS_ID = flags.DEFINE_integer(\"i\", -1, \"Process id\")\n 19_NUM_PROCESSES = flags.DEFINE_integer(\"n\", -1, \"Number of processes\")\n 20\n 21def replicated(x: jax.Array, devices: list[jax.Device]):\n 22  \"\"\"Return x replicated across the provided devices.\n 23\n 24  Note that replicated(x) doesn't actually move any data. It simply creates a\n 25  logically replicated array with x as the local replica.\n 26  \"\"\"\n 27  n = len(devices)\n 28  mesh = jax.make_mesh((n, ), (\"i\", ), devices=devices)\n 29  spec = jax.sharding.PartitionSpec(None)\n 30  sharding = jax.sharding.NamedSharding(mesh, spec)\n 31  shards = [\n 32      jax.device_put(x.addressable_shards[0].data, d) for d in devices\n 33      if d.process_index == jax.process_index()\n 34  ]\n 35  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)\n 36\n 37\n 38def sharded(x: jax.Array, devices: list[jax.Device]):\n 39  \"\"\"Return x sharded across the provided devices.\n 40\n 41  Note that sharded(x) doesn't actually move any data. It simply creates a\n 42  logically sharded array. x should have the same shape as the global array.\n 43  \"\"\"\n 44  n = len(devices)\n 45  mesh = jax.make_mesh((n, ), (\"i\", ), devices=devices)\n 46  spec = jax.sharding.PartitionSpec(\"i\")\n 47  sharding = jax.sharding.NamedSharding(mesh, spec)\n 48  m = sharding.addressable_devices_indices_map(x.shape)\n 49  shards = [jax.device_put(x[m[d]], d) for d in jax.local_devices()]\n 50  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)\n 51\n 52\n 53def main(_: Sequence[str]) -> None:\n 54  # Parse command line arguments and initialize multi-controller JAX.\n 55  jax.config.update(\"jax_enable_recoverability\", True)\n 56  jax.distributed.initialize(coordinator_address=\"localhost:8000\",\n 57                             process_id=_PROCESS_ID.value,\n 58                             num_processes=_NUM_PROCESSES.value,\n 59                             local_device_ids=[_PROCESS_ID.value],\n 60                             heartbeat_timeout_seconds=10)\n 61  print(f'{jax.devices()=}')\n 62  print(f'{jax.local_devices()=}')\n 63\n 64  # Initialize the model's weights.\n 65  keys = iter(jax.random.split(jax.random.key(seed=42), num=3))\n 66  weights = jax.random.normal(next(keys), shape=(1, ))\n 67\n 68  # We'll learn a trivial linear model: a*x.\n 69  def predict(weights, X):\n 70    return weights * X\n 71\n 72  # We'll use mean squared error loss.\n 73  def loss(weights, X, Y):\n 74    return jnp.mean((predict(weights, X) - Y)**2)\n 75\n 76  # Initialize the (noisy) training data with a=10.\n 77  X = jax.random.permutation(next(keys), jnp.arange(-300., 300.))\n 78  Y = 10 * X + jax.random.normal(next(keys), X.shape)\n 79\n 80  # Hyperparameters.\n 81  loss_and_grad = jax.jit(jax.value_and_grad(loss))\n 82  learning_rate = 1e-6\n 83  device_batch_size = 10\n 84\n 85  step = 0\n 86  while True:\n 87    try:\n 88      with live_devices(jax.devices()) as devices:\n 89        print(f'=== Running step {step} with live devices = {devices} ===')\n 90\n 91        # Replicate the model weights.\n 92        weights = replicated(weights, devices)\n 93\n 94        # Shard the batch.\n 95        batch_size = device_batch_size * len(devices)\n 96        start = (step * batch_size) % len(X)\n 97        stop = start + batch_size\n 98        X_batch = sharded(X[start:stop], devices)\n 99        Y_batch = sharded(Y[start:stop], devices)\n100\n101        # Compute gradients and update weights.\n102        l, grad = loss_and_grad(weights, X_batch, Y_batch)\n103        new_weights = jax.block_until_ready(weights - learning_rate * grad)\n104    except Exception as e:\n105      print(f'Step {step} failed: {e}')\n106    else:\n107      print(f'Step {step} succeeded: loss = {l}')\n108      step += 1\n109      weights = new_weights\n110\n111    time.sleep(1)\n112\n113\n114if __name__ == \"__main__\":\n115  app.run(main)\n```\n\nExample:\n```text\n55def send(x: jax.Array, from_device: jax.Device, to_device: jax.Device):\n56  \"\"\"Sends x from one device to another.\"\"\"\n57  assert isinstance(x, jax.Array)\n58  devices = [from_device, to_device]\n59  psum = lambda x: jax.lax.psum(x, \"i\")\n60  mesh = jax.make_mesh((2, ), (\"i\", ), devices=devices)\n61  spec = jax.sharding.PartitionSpec(None)\n62  x = replicated(x, [from_device, to_device])\n63  shard_map.shard_map(psum, mesh=mesh, in_specs=spec, out_specs=spec)(x)\n64\n65\n66def recv(x: jax.Array, from_device: jax.Device, to_device: jax.Device):\n67  \"\"\"Receives x from a matching send.\"\"\"\n68  assert isinstance(x, jax.Array)\n69  to_device = jax.local_devices()[0]\n70  devices = [from_device, to_device]\n71  psum = lambda x: jax.lax.psum(x, \"i\")\n72  mesh = jax.make_mesh((2, ), (\"i\", ), devices=devices)\n73  spec = jax.sharding.PartitionSpec(None)\n74  x = jnp.zeros_like(x)\n75  x = replicated(x, [from_device, to_device])\n76  return shard_map.shard_map(psum, mesh=mesh, in_specs=spec, out_specs=spec)(x)\n```\n\nExample:\n```text\n79def allgather(x: float, devices: list[jax.Device]) -> list[float]:\n80  \"\"\"Performs an AllGather across the provided devices.\"\"\"\n81  n = len(devices)\n82  mesh = jax.make_mesh((n, ), (\"i\", ), devices=devices)\n83  spec = jax.sharding.PartitionSpec('i')\n84  p = lambda x: jax.lax.all_gather(x, \"i\", tiled=True)\n85  f = jax.shard_map(p, mesh=mesh, in_specs=spec, out_specs=spec)\n86  return jax.block_until_ready(f(np.array([x] * len(devices)))).addressable_shards[0].data\n```\n\nExample:\n```text\n121  step = 0\n122  while True:\n123    try:\n124      with live_devices(jax.devices()) as devices:\n125        print(f'=== Running step {step} with live devices = {devices} ===')\n126\n127        # Handle recovering devices. A device is recovering if its step doesn't\n128        # match process 0's step. We assume process 0 never fails.\n129        print('all gathering steps...')\n130        steps = allgather(step, devices)\n131        print(f'{steps=}')\n132        recovering = [d for d, s in zip(devices, steps) if s != steps[0]]\n133        for d in recovering:\n134          # Process 0 sends weights and step to the recovering devices.\n135          if jax.process_index() == 0:\n136            print('sending...')\n137            send(weights, jax.devices()[0], d)\n138            send(jnp.array([step]), jax.devices()[0], d)\n139          elif d.process_index == jax.process_index():\n140            print('receiving...')\n141            weights = recv(weights, jax.devices()[0], d)\n142            step = recv(jnp.array([step]), jax.devices()[0], d)[0]\n143\n144        # Replicate the model weights.\n145        weights = replicated(weights, devices)\n146\n147        # Shard the batch.\n148        batch_size = device_batch_size * len(devices)\n149        start = (step * batch_size) % len(X)\n150        stop = start + batch_size\n151        X_batch = sharded(X[start:stop], devices)\n152        Y_batch = sharded(Y[start:stop], devices)\n153\n154        # Compute gradients and update weights.\n155        l, grad = loss_and_grad(weights, X_batch, Y_batch)\n156        new_weights = jax.block_until_ready(weights - learning_rate * grad)\n157    except Exception as e:\n158      print(f'Step {step} failed: {e}')\n159    else:\n160      print(f'Step {step} succeeded: loss = {l}')\n161      step += 1\n162      weights = new_weights\n163\n164    time.sleep(1)\n```\n\nExample:\n```text\n1import os\n  2os.environ['XLA_FLAGS'] = ' '.join([\n  3    '--xla_gpu_nccl_terminate_on_error=false',\n  4    '--xla_gpu_nccl_async_execution=true',\n  5    '--xla_gpu_nccl_blocking_communicators=false',\n  6])\n  7os.environ['XLA_PYTHON_CLIENT_ABORT_COLLECTIVES_ON_FAILURE'] = '1'\n  8os.environ['XLA_PYTHON_CLIENT_USE_TFRT_GPU_CLIENT'] = '1'\n  9\n 10from absl import app\n 11from absl import flags\n 12from collections.abc import Sequence\n 13from jax.experimental.multihost_utils import live_devices\n 14from jax.experimental import shard_map\n 15import jax\n 16import jax.numpy as jnp\n 17import numpy as np\n 18import time\n 19\n 20_PROCESS_ID = flags.DEFINE_integer(\"i\", -1, \"Process id\")\n 21_NUM_PROCESSES = flags.DEFINE_integer(\"n\", -1, \"Number of processes\")\n 22\n 23def replicated(x: jax.Array, devices: list[jax.Device]):\n 24  \"\"\"Return x replicated across the provided devices.\n 25\n 26  Note that replicated(x) doesn't actually move any data. It simply creates a\n 27  logically replicated array with x as the local replica.\n 28  \"\"\"\n 29  n = len(devices)\n 30  mesh = jax.make_mesh((n, ), (\"i\", ), devices=devices)\n 31  spec = jax.sharding.PartitionSpec(None)\n 32  sharding = jax.sharding.NamedSharding(mesh, spec)\n 33  shards = [\n 34      jax.device_put(x.addressable_shards[0].data, d) for d in devices\n 35      if d.process_index == jax.process_index()\n 36  ]\n 37  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)\n 38\n 39\n 40def sharded(x: jax.Array, devices: list[jax.Device]):\n 41  \"\"\"Return x sharded across the provided devices.\n 42\n 43  Note that sharded(x) doesn't actually move any data. It simply creates a\n 44  logically sharded array. x should have the same shape as the global array.\n 45  \"\"\"\n 46  n = len(devices)\n 47  mesh = jax.make_mesh((n, ), (\"i\", ), devices=devices)\n 48  spec = jax.sharding.PartitionSpec(\"i\")\n 49  sharding = jax.sharding.NamedSharding(mesh, spec)\n 50  m = sharding.addressable_devices_indices_map(x.shape)\n 51  shards = [jax.device_put(x[m[d]], d) for d in jax.local_devices()]\n 52  return jax.make_array_from_single_device_arrays(x.shape, sharding, shards)\n 53\n 54\n 55def send(x: jax.Array, from_device: jax.Device, to_device: jax.Device):\n 56  \"\"\"Sends x from one device to another.\"\"\"\n 57  assert isinstance(x, jax.Array)\n 58  devices = [from_device, to_device]\n 59  psum = lambda x: jax.lax.psum(x, \"i\")\n 60  mesh = jax.make_mesh((2, ), (\"i\", ), devices=devices)\n 61  spec = jax.sharding.PartitionSpec(None)\n 62  x = replicated(x, [from_device, to_device])\n 63  shard_map.shard_map(psum, mesh=mesh, in_specs=spec, out_specs=spec)(x)\n 64\n 65\n 66def recv(x: jax.Array, from_device: jax.Device, to_device: jax.Device):\n 67  \"\"\"Receives x from a matching send.\"\"\"\n 68  assert isinstance(x, jax.Array)\n 69  to_device = jax.local_devices()[0]\n 70  devices = [from_device, to_device]\n 71  psum = lambda x: jax.lax.psum(x, \"i\")\n 72  mesh = jax.make_mesh((2, ), (\"i\", ), devices=devices)\n 73  spec = jax.sharding.PartitionSpec(None)\n 74  x = jnp.zeros_like(x)\n 75  x = replicated(x, [from_device, to_device])\n 76  return shard_map.shard_map(psum, mesh=mesh, in_specs=spec, out_specs=spec)(x)\n 77\n 78\n 79def allgather(x: float, devices: list[jax.Device]) -> list[float]:\n 80  \"\"\"Performs an AllGather across the provided devices.\"\"\"\n 81  n = len(devices)\n 82  mesh = jax.make_mesh((n, ), (\"i\", ), devices=devices)\n 83  spec = jax.sharding.PartitionSpec('i')\n 84  p = lambda x: jax.lax.all_gather(x, \"i\", tiled=True)\n 85  f = jax.shard_map(p, mesh=mesh, in_specs=spec, out_specs=spec)\n 86  return jax.block_until_ready(f(np.array([x] * len(devices)))).addressable_shards[0].data\n 87\n 88\n 89def main(_: Sequence[str]) -> None:\n 90  # Parse command line arguments and initialize multi-controller JAX.\n 91  jax.config.update(\"jax_enable_recoverability\", True)\n 92  jax.distributed.initialize(coordinator_address=\"localhost:8000\",\n 93                             process_id=_PROCESS_ID.value,\n 94                             num_processes=_NUM_PROCESSES.value,\n 95                             local_device_ids=[_PROCESS_ID.value],\n 96                             heartbeat_timeout_seconds=10)\n 97  print(f'{jax.devices()=}')\n 98  print(f'{jax.local_devices()=}')\n 99\n100  # Initialize the model's weights.\n101  keys = iter(jax.random.split(jax.random.key(seed=42), num=3))\n102  weights = jax.random.normal(next(keys), shape=(1, ))\n103\n104  # We'll learn a trivial linear model: a*x.\n105  def predict(weights, X):\n106    return weights * X\n107\n108  # We'll use mean squared error loss.\n109  def loss(weights, X, Y):\n110    return jnp.mean((predict(weights, X) - Y)**2)\n111\n112  # Initialize the (noisy) training data with a=10.\n113  X = jax.random.permutation(next(keys), jnp.arange(-300., 300.))\n114  Y = 10 * X + jax.random.normal(next(keys), X.shape)\n115\n116  # Hyperparameters.\n117  loss_and_grad = jax.jit(jax.value_and_grad(loss))\n118  learning_rate = 1e-6\n119  device_batch_size = 10\n120\n121  step = 0\n122  while True:\n123    try:\n124      with live_devices(jax.devices()) as devices:\n125        print(f'=== Running step {step} with live devices = {devices} ===')\n126\n127        # Handle recovering devices. A device is recovering if its step doesn't\n128        # match process 0's step. We assume process 0 never fails.\n129        print('all gathering steps...')\n130        steps = allgather(step, devices)\n131        print(f'{steps=}')\n132        recovering = [d for d, s in zip(devices, steps) if s != steps[0]]\n133        for d in recovering:\n134          # Process 0 sends weights and step to the recovering devices.\n135          if jax.process_index() == 0:\n136            print('sending...')\n137            send(weights, jax.devices()[0], d)\n138            send(jnp.array([step]), jax.devices()[0], d)\n139          elif d.process_index == jax.process_index():\n140            print('receiving...')\n141            weights = recv(weights, jax.devices()[0], d)\n142            step = recv(jnp.array([step]), jax.devices()[0], d)[0]\n143\n144        # Replicate the model weights.\n145        weights = replicated(weights, devices)\n146\n147        # Shard the batch.\n148        batch_size = device_batch_size * len(devices)\n149        start = (step * batch_size) % len(X)\n150        stop = start + batch_size\n151        X_batch = sharded(X[start:stop], devices)\n152        Y_batch = sharded(Y[start:stop], devices)\n153\n154        # Compute gradients and update weights.\n155        l, grad = loss_and_grad(weights, X_batch, Y_batch)\n156        new_weights = jax.block_until_ready(weights - learning_rate * grad)\n157    except Exception as e:\n158      print(f'Step {step} failed: {e}')\n159    else:\n160      print(f'Step {step} succeeded: loss = {l}')\n161      step += 1\n162      weights = new_weights\n163\n164    time.sleep(1)\n165\n166\n167if __name__ == \"__main__\":\n168  app.run(main)\n```\n\nExample:\n```text\nstep = 0\nwhile True:\n    # Get the devices on all live processes.\n    procs = jax.live_processes()\n    devices = [d for d in jax.devices() if d.process_index in procs]\n\n    # Shard array x over these devices.\n    mesh = jax.make_mesh((len(devices),), (\"i\",), devices=devices)\n    spec = jax.sharding.PartitionSpec(\"i\")\n    sharding = jax.sharding.NamedSharding(mesh, spec)\n    x = jax.make_array_from_process_local_data(sharding, np.ones(1))\n\n    # Try to perform a jnp.sum.\n    try:\n        print(jnp.sum(x))\n    except:\n        # jnp.sum failed.\n        pass\n    else:\n        # jnp.sum succeeded.\n        step += 1\n```\n\nExample:\n```text\n# Get the set of live processes before the code block.\nprocs_before = jax.live_processes()\n\n# Execute the code block.\n...\n\n# Get the set of live processes after the code block\nprocs_after = jax.live_processes()\nif procs_before == procs_after:\n    # The code block executed successfully on all processes in\n    # procs_before.\n    pass\nelse:\n    # The code block did not execute successfully. All processes will\n    # agree it failed.\n    pass\n```\n\nExample:\n```text\ntry:\n  with live_devices() as devices:\n    pass # A\nexcept Exception as e:\n  pass # B\nelse:\n  pass # C\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.830Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":26,"totalLines":842,"estimatedTokens":7164}}114{"id":"doc-jax_internals_the_jaxpr_language_jax_documentati-83b60850","source":"documentation","title":"JAX internals: The jaxpr language — JAX documentation","url":"https://docs.jax.dev/en/latest/jaxpr.html","text":"Example:\n```text\njaxpr ::=\n  { lambda <binder> , ... .\n    let <eqn>\n        ...\n    in ( <atom> , ... ) }\n\nbinder ::= <var>:<array_type>\nvar ::= a | b | c | ...\natom ::= <var> | <literal>\nliteral ::= <int32> | <int64> | <float32> | <float64>\n\neqn ::= <binder> , ... = <primitive> [ <params> ] <atom> , ...\n```\n\nExample:\n```text\njaxpr ::= { lambda Var* ; Var+.\n            let Eqn*\n            in  [Expr+] }\n```\n\nExample:\n```text\nEqn  ::= let Var+ = Primitive [ Param* ] Expr+\n```\n\nExample:\n```text\nPrimitive := add | sub | sin | mul | ...\n```\n\nExample:\n```text\nfrom jax import make_jaxpr\nimport jax.numpy as jnp\n\ndef func1(first, second):\n   temp = first + jnp.sin(second) * 3.\n   return jnp.sum(temp)\n\nprint(make_jaxpr(func1)(jnp.zeros(8), jnp.ones(8)))\n```\n\nExample:\n```text\n{ lambda ; a:f32[8] b:f32[8]. let\n    c:f32[8] = sin b\n    d:f32[8] = mul c 3.0:f32[]\n    e:f32[8] = add a d\n    f:f32[] = reduce_sum[axes=(0,) out_sharding=None] e\n  in (f,) }\n```\n\nExample:\n```text\ndef func2(inner, first, second):\n  temp = first + inner(second) * 3.\n  return jnp.sum(temp)\n\ndef inner(second):\n  if second.shape[0] > 4:\n    return jnp.sin(second)\n  else:\n    assert False\n\ndef func3(first, second):\n  return func2(inner, first, second)\n\nprint(make_jaxpr(func3)(jnp.zeros(8), jnp.ones(8)))\n```\n\nExample:\n```text\ndef func4(arg):  # The `arg` is a pair.\n  temp = arg[0] + jnp.sin(arg[1]) * 3.\n  return jnp.sum(temp)\n\nprint(make_jaxpr(func4)((jnp.zeros(8), jnp.ones(8))))\n```\n\nExample:\n```text\nlax.switch(index: int, branches: Sequence[A -> B], operand: A) -> B\n\nlax.cond(pred: bool, true_body: A -> B, false_body: A -> B, operand: A) -> B\n```\n\nExample:\n```text\nfrom jax import lax\n\ndef one_of_three(index, arg):\n  return lax.switch(index, [lambda x: x + 1.,\n                            lambda x: x - 2.,\n                            lambda x: x + 3.],\n                    arg)\n\nprint(make_jaxpr(one_of_three)(1, 5.))\n```\n\nExample:\n```text\n{ lambda ; a:i32[] b:f32[]. let\n    c:i32[] = convert_element_type[new_dtype=int32 weak_type=False] a\n    d:i32[] = clamp 0:i32[] c 2:i32[]\n    e:f32[] = cond[\n      branches=(\n        { lambda ; f:f32[]. let g:f32[] = add f 1.0:f32[] in (g,) }\n        { lambda ; h:f32[]. let i:f32[] = sub h 2.0:f32[] in (i,) }\n        { lambda ; j:f32[]. let k:f32[] = add j 3.0:f32[] in (k,) }\n      )\n    ] d b\n  in (e,) }\n```\n\nExample:\n```text\nfrom jax import lax\n\ndef func7(arg):\n  return lax.cond(arg >= 0.,\n                  lambda xtrue: xtrue + 3.,\n                  lambda xfalse: xfalse - 3.,\n                  arg)\n\nprint(make_jaxpr(func7)(5.))\n```\n\nExample:\n```text\n{ lambda ; a:f32[]. let\n    b:bool[] = ge a 0.0:f32[]\n    c:i32[] = convert_element_type[new_dtype=int32 weak_type=False] b\n    d:f32[] = cond[\n      branches=(\n        { lambda ; e:f32[]. let f:f32[] = sub e 3.0:f32[] in (f,) }\n        { lambda ; g:f32[]. let h:f32[] = add g 3.0:f32[] in (h,) }\n      )\n    ] c a\n  in (d,) }\n```\n\nExample:\n```text\ndef func8(arg1, arg2):  # Where `arg2` is a pair.\n  return lax.cond(arg1 >= 0.,\n                  lambda xtrue: xtrue[0],\n                  lambda xfalse: jnp.array([1]) + xfalse[1],\n                  arg2)\n\nprint(make_jaxpr(func8)(5., (jnp.zeros(1), 2.)))\n```\n\nExample:\n```text\n{ lambda a:i32[1]; b:f32[] c:f32[1] d:f32[]. let\n    e:bool[] = ge b 0.0:f32[]\n    f:i32[] = convert_element_type[new_dtype=int32 weak_type=False] e\n    g:f32[1] = cond[\n      branches=(\n        { lambda ; h:i32[1] i:f32[1] j:f32[]. let\n            k:f32[1] = convert_element_type[new_dtype=float32 weak_type=True] h\n            l:f32[1] = add k j\n          in (l,) }\n        { lambda ; m:i32[1] n:f32[1] o:f32[]. let  in (n,) }\n      )\n    ] f a c d\n  in (g,) }\n```\n\nExample:\n```text\nlax.while_loop(cond_fun: (C -> bool), body_fun: (C -> C), init: C) -> C\nlax.fori_loop(start: int, end: int, body: (int -> C -> C), init: C) -> C\n```\n\nExample:\n```text\nimport numpy as np\n\ndef func10(arg, n):\n  ones = jnp.ones(arg.shape)  # A constant.\n  return lax.fori_loop(0, n,\n                       lambda i, carry: carry + ones * 3. + arg,\n                       arg + ones)\n\nprint(make_jaxpr(func10)(np.ones(16), 5))\n```\n\nExample:\n```text\n{ lambda ; a:f32[16] b:i32[]. let\n    c:f32[16] = broadcast_in_dim 1.0:f32[]\n    d:f32[16] = add a c\n    _:i32[] _:i32[] e:f32[16] = while[\n      body_jaxpr={ lambda ; f:f32[16] g:f32[16] h:i32[] i:i32[] j:f32[16]. let\n          k:i32[] = add h 1:i32[]\n          l:f32[16] = mul f 3.0:f32[]\n          m:f32[16] = add j l\n          n:f32[16] = add m g\n        in (k, i, n) }\n      body_nconsts=2\n      cond_jaxpr={ lambda ; o:i32[] p:i32[] q:f32[16]. let\n          r:bool[] = lt o p\n        in (r,) }\n      cond_nconsts=0\n    ] c a 0:i32[] b d\n  in (e,) }\n```\n\nExample:\n```text\nlax.scan(body_fun: (C -> A -> (C, B)), init_carry: C, in_arr: Array[A]) -> (C, Array[B])\n```\n\nExample:\n```text\ndef func11(arr, extra):\n  ones = jnp.ones(arr.shape)  #  A constant\n  def body(carry, aelems):\n    # carry: running dot-product of the two arrays\n    # aelems: a pair with corresponding elements from the two arrays\n    ae1, ae2 = aelems\n    return (carry + ae1 * ae2 + extra, carry)\n  return lax.scan(body, 0., (arr, ones))\n\nprint(make_jaxpr(func11)(np.ones(16), 5.))\n```\n\nExample:\n```text\n{ lambda ; a:f32[16] b:f32[]. let\n    c:f32[16] = broadcast_in_dim 1.0:f32[]\n    d:f32[] e:f32[16] = scan[\n      ft_in=((None,), (None,), (None, None))\n      ft_out=((None,), (RightsOnly(Right(None)),))\n      jaxpr={ lambda ; f:f32[] g:f32[] h:f32[] i:f32[]. let\n          j:f32[] = mul h i\n          k:f32[] = convert_element_type[new_dtype=float32 weak_type=False] g\n          l:f32[] = add k j\n          m:f32[] = convert_element_type[new_dtype=float32 weak_type=False] f\n          n:f32[] = add l m\n        in (n, g) }\n      length=16\n      reverse=False\n      unroll=1\n    ] b 0.0:f32[] a c\n  in (d, e) }\n```\n\nExample:\n```text\nfrom jax import jit\n\ndef func12(arg):\n  @jit\n  def inner(x):\n    return x + arg * jnp.ones(1)  # Include a constant in the inner function.\n  return arg + inner(arg - 2.)\n\nprint(make_jaxpr(func12)(1.))\n```\n\nExample:\n```text\n{ lambda ; a:f32[]. let\n    b:f32[] = sub a 2.0:f32[]\n    c:f32[1] = jit[\n      name=inner\n      jaxpr={ lambda ; a:f32[] b:f32[]. let\n          d:f32[1] = broadcast_in_dim 1.0:f32[]\n          e:f32[] = convert_element_type[new_dtype=float32 weak_type=False] a\n          f:f32[1] = mul e d\n          g:f32[] = convert_element_type[new_dtype=float32 weak_type=False] b\n          c:f32[1] = add g f\n        in (c,) }\n    ] a b\n    h:f32[] = convert_element_type[new_dtype=float32 weak_type=False] a\n    i:f32[1] = add h c\n  in (i,) }\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.832Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":23,"totalLines":285,"estimatedTokens":1659}}115{"id":"doc-jax_internals_primitives_jax_documentation-07b5295b","source":"documentation","title":"JAX Internals: primitives — JAX documentation","url":"https://docs.jax.dev/en/latest/jax-primitives.html","text":"Example:\n```text\nfrom jax._src.lax import lax\nfrom jax._src import api\n\ndef multiply_add_lax(x, y, z):\n  \"\"\"Implementation of multiply-add using the `jax.lax` primitives.\"\"\"\n  return lax.add(lax.mul(x, y), z)\n\n\ndef square_add_lax(a, b):\n  \"\"\"A square-add function using the newly defined multiply-add.\"\"\"\n  return multiply_add_lax(a, a, b)\n\nprint(\"square_add_lax = \", square_add_lax(2., 10.))\n# Differentiate w.r.t. the first argument\nprint(\"grad(square_add_lax) = \", api.grad(square_add_lax, argnums=0)(2.0, 10.))\n```\n\nExample:\n```text\nsquare_add_lax =  14.0\ngrad(square_add_lax) =  4.0\n```\n\nExample:\n```text\n#@title Helper functions (execute this cell)\nimport functools\nimport traceback\n\n_indentation = 0\ndef _trace(msg=None):\n    \"\"\"Print a message at current indentation.\"\"\"\n    if msg is not None:\n        print(\"  \" * _indentation + msg)\n\ndef _trace_indent(msg=None):\n    \"\"\"Print a message and then indent the rest.\"\"\"\n    global _indentation\n    _trace(msg)\n    _indentation = 1 + _indentation\n\ndef _trace_unindent(msg=None):\n    \"\"\"Unindent then print a message.\"\"\"\n    global _indentation\n    _indentation = _indentation - 1\n    _trace(msg)\n\ndef trace(name):\n  \"\"\"A decorator for functions to trace arguments and results.\"\"\"\n\n  def trace_func(func):\n    def pp(v):\n        \"\"\"Print certain values more succinctly\"\"\"\n        vtype = str(type(v))\n        if \"jax._src.xla_bridge._JaxComputationBuilder\" in vtype:\n            return \"<JaxComputationBuilder>\"\n        elif \"jaxlib._jax_.XlaOp\" in vtype:\n            return \"<XlaOp at 0x{:x}>\".format(id(v))\n        elif (\"partial_eval.JaxprTracer\" in vtype or\n              \"batching.BatchTracer\" in vtype or\n              \"ad.JVPTracer\" in vtype):\n            return \"Traced<{}>\".format(v.aval)\n        elif isinstance(v, tuple):\n            return \"({})\".format(pp_values(v))\n        else:\n            return str(v)\n    def pp_values(args):\n        return \", \".join([pp(arg) for arg in args])\n\n    @functools.wraps(func)\n    def func_wrapper(*args):\n      _trace_indent(\"call {}({})\".format(name, pp_values(args)))\n      res = func(*args)\n      _trace_unindent(\"|<- {} = {}\".format(name, pp(res)))\n      return res\n\n    return func_wrapper\n\n  return trace_func\n\nclass expectNotImplementedError(object):\n  \"\"\"Context manager to check for NotImplementedError.\"\"\"\n  def __enter__(self): pass\n  def __exit__(self, type, value, tb):\n    global _indentation\n    _indentation = 0\n    if type is NotImplementedError:\n      print(\"\\nFound expected exception:\")\n      traceback.print_exc(limit=3)\n      return True\n    elif type is None:  # No exception\n      assert False, \"Expected NotImplementedError\"\n    else:\n      return False\n```\n\nExample:\n```text\nimport jax.numpy as jnp\nimport numpy as np\n\n@trace(\"multiply_add_numpy\")\ndef multiply_add_numpy(x, y, z):\n    return jnp.add(jnp.multiply(x, y), z)\n\n@trace(\"square_add_numpy\")\ndef square_add_numpy(a, b):\n    return multiply_add_numpy(a, a, b)\n\nprint(\"\\nNormal evaluation:\")\nprint(\"square_add_numpy = \", square_add_numpy(2., 10.))\nprint(\"\\nGradient evaluation:\")\nprint(\"grad(square_add_numpy) = \", api.grad(square_add_numpy)(2.0, 10.))\n```\n\nExample:\n```text\nNormal evaluation:\ncall square_add_numpy(2.0, 10.0)\n  call multiply_add_numpy(2.0, 2.0, 10.0)\n  |<- multiply_add_numpy = 14.0\n|<- square_add_numpy = 14.0\nsquare_add_numpy =  14.0\n\nGradient evaluation:\ncall square_add_numpy(GradTracer(primal=2.0, typeof(tangent)=f32[]), 10.0)\n  call multiply_add_numpy(GradTracer(primal=2.0, typeof(tangent)=f32[]), GradTracer(primal=2.0, typeof(tangent)=f32[]), 10.0)\n  |<- multiply_add_numpy = GradTracer(primal=14.0, typeof(tangent)=f32[])\n|<- square_add_numpy = GradTracer(primal=14.0, typeof(tangent)=f32[])\ngrad(square_add_numpy) =  4.0\n```\n\nExample:\n```text\nfrom jax.extend import core\n\nmultiply_add_p = core.Primitive(\"multiply_add\")  # Create the primitive\n\n@trace(\"multiply_add_prim\")\ndef multiply_add_prim(x, y, z):\n  \"\"\"The JAX-traceable way to use the JAX primitive.\n\n  Note that the traced arguments must be passed as positional arguments\n  to `bind`.\n  \"\"\"\n  return multiply_add_p.bind(x, y, z)\n\n@trace(\"square_add_prim\")\ndef square_add_prim(a, b):\n  \"\"\"A square-add function implemented using the new JAX-primitive.\"\"\"\n  return multiply_add_prim(a, a, b)\n```\n\nExample:\n```text\nwith expectNotImplementedError():\n  square_add_prim(2., 10.)\n```\n\nExample:\n```text\ncall square_add_prim(2.0, 10.0)\n  call multiply_add_prim(2.0, 2.0, 10.0)\n\nFound expected exception:\n```\n\nExample:\n```text\nTraceback (most recent call last):\n  File \"/tmp/ipykernel_2054/2844449444.py\", line 2, in <module>\n    square_add_prim(2., 10.)\n  File \"/tmp/ipykernel_2054/3025987661.py\", line 48, in func_wrapper\n    res = func(*args)\n          ^^^^^^^^^^^\n  File \"/tmp/ipykernel_2054/3275395289.py\", line 17, in square_add_prim\n    return multiply_add_prim(a, a, b)\n           ^^^^^^^^^^^^^^^^^^^^^^^^^^\nNotImplementedError: Evaluation rule for 'multiply_add' not implemented\n```\n\nExample:\n```text\n@trace(\"multiply_add_impl\")\ndef multiply_add_impl(x, y, z):\n  \"\"\"Concrete implementation of the primitive.\n\n  This function does not need to be JAX traceable.\n\n  Args:\n    x, y, z: The concrete arguments of the primitive. Will only be called with\n      concrete values.\n\n  Returns:\n    the concrete result of the primitive.\n  \"\"\"\n  # Note: you can use the ordinary (non-JAX) NumPy, which is not JAX-traceable.\n  return np.add(np.multiply(x, y), z)\n\n# Now, register the primal implementation with JAX:\nmultiply_add_p.def_impl(multiply_add_impl)\n```\n\nExample:\n```text\n<function __main__.multiply_add_impl(x, y, z)>\n```\n\nExample:\n```text\nassert square_add_prim(2., 10.) == 14.\n```\n\nExample:\n```text\ncall square_add_prim(2.0, 10.0)\n  call multiply_add_prim(2.0, 2.0, 10.0)\n    call multiply_add_impl(2.0, 2.0, 10.0)\n    |<- multiply_add_impl = 14.0\n  |<- multiply_add_prim = 14.0\n|<- square_add_prim = 14.0\n```\n\nExample:\n```text\nwith expectNotImplementedError():\n  api.jit(square_add_prim)(2., 10.)\n```\n\nExample:\n```text\ncall square_add_prim(JitTracer(~float32[]), JitTracer(~float32[]))\n  call multiply_add_prim(JitTracer(~float32[]), JitTracer(~float32[]), JitTracer(~float32[]))\n\nFound expected exception:\n```\n\nExample:\n```text\nTraceback (most recent call last):\n  File \"/tmp/ipykernel_2054/1813425700.py\", line 2, in <module>\n    api.jit(square_add_prim)(2., 10.)\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/traceback_util.py\", line 194, in reraise_with_filtered_traceback\n    return fun(*args, **kwargs)  # pyrefly: ignore[not-callable]\n           ^^^^^^^^^^^^^^^^^^^^\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/pjit.py\", line 263, in cache_miss\n    p, args_flat = _infer_params(fun, jit_info, args, kwargs)\n                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nNotImplementedError: Abstract evaluation for 'multiply_add' not implemented\n```\n\nExample:\n```text\nfrom jax import core\n\n@trace(\"multiply_add_abstract_eval\")\ndef multiply_add_abstract_eval(xs, ys, zs):\n  \"\"\"Abstract evaluation of the primitive.\n\n  This function does not need to be JAX traceable. It will be invoked with\n  abstractions of the actual arguments\n\n  Args:\n    xs, ys, zs: Abstractions of the arguments.\n\n  Result:\n    a ShapedArray for the result of the primitive.\n  \"\"\"\n  assert xs.shape == ys.shape\n  assert xs.shape == zs.shape\n  return core.ShapedArray(xs.shape, xs.dtype)\n\n# Now, register the abstract evaluation with JAX:\nmultiply_add_p.def_abstract_eval(multiply_add_abstract_eval)\n```\n\nExample:\n```text\n<function __main__.multiply_add_abstract_eval(xs, ys, zs)>\n```\n\nExample:\n```text\ncall square_add_prim(JitTracer(~float32[]), JitTracer(~float32[]))\n  call multiply_add_prim(JitTracer(~float32[]), JitTracer(~float32[]), JitTracer(~float32[]))\n    call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])\n    |<- multiply_add_abstract_eval = float32[]\n  |<- multiply_add_prim = JitTracer(float32[])\n|<- square_add_prim = JitTracer(float32[])\n\nFound expected exception:\n```\n\nExample:\n```text\nTraceback (most recent call last):\n  File \"/tmp/ipykernel_2054/3025987661.py\", line 48, in func_wrapper\n    res = func(*args)\n  File \"/tmp/ipykernel_2054/3275395289.py\", line 17, in square_add_prim\n    return multiply_add_prim(a, a, b)\n  File \"/tmp/ipykernel_2054/3025987661.py\", line 48, in func_wrapper\n    res = func(*args)\njax._src.source_info_util.JaxStackTraceBeforeTransformation: NotImplementedError: MLIR translation rule for primitive 'multiply_add' not found for platform cpu\n\nThe preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.\n\n--------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n  File \"/tmp/ipykernel_2054/1813425700.py\", line 2, in <module>\n    api.jit(square_add_prim)(2., 10.)\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/traceback_util.py\", line 194, in reraise_with_filtered_traceback\n    return fun(*args, **kwargs)  # pyrefly: ignore[not-callable]\n           ^^^^^^^^^^^^^^^^^^^^\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/pjit.py\", line 265, in cache_miss\n    executable, pgle_profiler, const_args) = _run_python_pjit(\n                                             ^^^^^^^^^^^^^^^^^\nNotImplementedError: MLIR translation rule for primitive 'multiply_add' not found for platform cpu\n```\n\nExample:\n```text\nfrom jax._src.lib.mlir.dialects import hlo\n\n@trace(\"multiply_add_lowering\")\ndef multiply_add_lowering(ctx, xc, yc, zc):\n  \"\"\"The compilation to XLA of the primitive.\n\n  Given an mlir.ir.Value for each argument, return the mlir.ir.Values for\n  the results of the function.\n\n  Does not need to be a JAX-traceable function.\n  \"\"\"\n  return [hlo.AddOp(hlo.MulOp(xc, yc), zc).result]\n\n# Now, register the lowering rule with JAX.\n# For GPU, refer to the https://docs.jax.dev/en/latest/Custom_Operation_for_GPUs.html\nfrom jax.interpreters import mlir\n\nmlir.register_lowering(multiply_add_p, multiply_add_lowering, platform='cpu')\n```\n\nExample:\n```text\nassert api.jit(lambda x, y: square_add_prim(x, y))(2., 10.) == 14.\n```\n\nExample:\n```text\ncall square_add_prim(JitTracer(~float32[]), JitTracer(~float32[]))\n  call multiply_add_prim(JitTracer(~float32[]), JitTracer(~float32[]), JitTracer(~float32[]))\n    call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])\n    |<- multiply_add_abstract_eval = float32[]\n  |<- multiply_add_prim = JitTracer(float32[])\n|<- square_add_prim = JitTracer(float32[])\ncall multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x7a00db6720d0>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x7a00d82584a0>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x7a00d82584e0>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x7a00d830f300>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x7a00db66c940>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x7a00d830e7e0>, all_default_mem_kind=True, lowering_cache={}, cached_primitive_lowerings={}, sharding_attr_cache={}, aval_to_ir_types_cache={ShapedArray(float32[], weak_type=True): RankedTensorType(tensor<f32>), ShapedArray(float32[]): RankedTensorType(tensor<f32>), ShapedArray(int32[]): RankedTensorType(tensor<i32>)}, pallas_lowering_cache={}, pallas_collective_id_mapping=CollectiveIdMapping(auto={}, manual={}, all_ids=set()), traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x7a00d8247eb0>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x7a00d830ff40>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), remove_size_one_mesh_axis=False, xla_metadata=None), platforms=None), BlockArgument(<block argument> of type 'tensor<f32>' at index: 0), BlockArgument(<block argument> of type 'tensor<f32>' at index: 1), BlockArgument(<block argument> of type 'tensor<f32>' at index: 2))\n|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x7a00d8c49eb0>]\n```\n\nExample:\n```text\nassert api.jit(lambda x, y: square_add_prim(x, y),\n               static_argnums=1)(2., 10.) == 14.\n```\n\nExample:\n```text\ncall square_add_prim(JitTracer(~float32[]), 10.0)\n  call multiply_add_prim(JitTracer(~float32[]), JitTracer(~float32[]), 10.0)\n    call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])\n    |<- multiply_add_abstract_eval = float32[]\n  |<- multiply_add_prim = JitTracer(float32[])\n|<- square_add_prim = JitTracer(float32[])\ncall multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x7a00d8247ed0>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x7a00d8258ea0>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x7a00d8258ee0>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x7a00d8260e40>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x7a00db66c940>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x7a00d8260e90>, all_default_mem_kind=True, lowering_cache={}, cached_primitive_lowerings={}, sharding_attr_cache={}, aval_to_ir_types_cache={ShapedArray(float32[], weak_type=True): RankedTensorType(tensor<f32>), ShapedArray(float32[]): RankedTensorType(tensor<f32>), ShapedArray(int32[]): RankedTensorType(tensor<i32>)}, pallas_lowering_cache={}, pallas_collective_id_mapping=CollectiveIdMapping(auto={}, manual={}, all_ids=set()), traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x7a00d825cb30>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x7a00d82617b0>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), remove_size_one_mesh_axis=False, xla_metadata=None), platforms=None), BlockArgument(<block argument> of type 'tensor<f32>' at index: 0), BlockArgument(<block argument> of type 'tensor<f32>' at index: 1), BlockArgument(<block argument> of type 'tensor<f32>' at index: 2))\n|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x7a00d8b7fe30>]\n```\n\nExample:\n```text\n# The second argument is set to `(2., 10.)` values where you\n# evaluate the Jacobian, and the third argument `(1., 1.)`\n# contains the values of the tangents for the arguments.\nwith expectNotImplementedError():\n  api.jvp(square_add_prim, (2., 10.), (1., 1.))\n```\n\nExample:\n```text\ncall square_add_prim(Traced<~float32[]>, Traced<~float32[]>)\n  call multiply_add_prim(Traced<~float32[]>, Traced<~float32[]>, Traced<~float32[]>)\n\nFound expected exception:\n```\n\nExample:\n```text\nTraceback (most recent call last):\n  File \"/tmp/ipykernel_2054/459539105.py\", line 5, in <module>\n    api.jvp(square_add_prim, (2., 10.), (1., 1.))\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/traceback_util.py\", line 194, in reraise_with_filtered_traceback\n    return fun(*args, **kwargs)  # pyrefly: ignore[not-callable]\n           ^^^^^^^^^^^^^^^^^^^^\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/api.py\", line 1450, in jvp\n    return _jvp(fun, primals, tangents, has_aux=has_aux)\n           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nNotImplementedError: Differentiation rule for 'multiply_add' not implemented\n```\n\nExample:\n```text\nfrom jax.interpreters import ad\n\n@trace(\"multiply_add_value_and_jvp\")\ndef multiply_add_value_and_jvp(arg_values, arg_tangents):\n  \"\"\"Evaluates the primal output and the tangents (Jacobian-vector product).\n\n  Given values of the arguments and perturbation of the arguments (tangents),\n  compute the output of the primitive and the perturbation of the output.\n\n  This method must be JAX-traceable. JAX may invoke it with abstract values\n  for the arguments and tangents.\n\n  Args:\n    arg_values: A tuple of arguments\n    arg_tangents: A tuple with the tangents of the arguments. The tuple has\n      the same length as the arg_values. Some of the tangents may also be the\n      special value `ad.Zero` to specify a zero tangent\n\n  Returns:\n     A pair of the primal output and the tangent.\n  \"\"\"\n  x, y, z = arg_values\n  xt, yt, zt = arg_tangents\n  _trace(\"Primal evaluation:\")\n  # Now, you have a JAX-traceable computation of the output.\n  # Normally, you can use the multiply add (`ma`) primitive itself to compute the primal output.\n  primal_out = multiply_add_prim(x, y, z)\n\n  _trace(\"Tangent evaluation:\")\n  # You must use a JAX-traceable way to compute the tangent. It turns out that\n  # the output tangent can be computed as (xt * y + x * yt + zt),\n  # which you can implement in a JAX-traceable way using the same \"multiply_add_prim\" primitive.\n\n  # You do need to deal specially with `Zero`. Here, you just turn it into a\n  # proper tensor of 0s (of the same shape as 'x').\n  # An alternative would be to check for `Zero` and perform algebraic\n  # simplification of the output tangent computation.\n  def make_zero(tan):\n    return lax.full_like(x, 0) if type(tan) is ad.Zero else tan\n\n  output_tangent = multiply_add_prim(make_zero(xt), y, multiply_add_prim(x, make_zero(yt), make_zero(zt)))\n  return (primal_out, output_tangent)\n\n# Register the forward differentiation rule with JAX:\nad.primitive_jvps[multiply_add_p] = multiply_add_value_and_jvp\n```\n\nExample:\n```text\n# Tangent is: xt*y + x*yt + zt = 1.*2. + 2.*1. + 1. = 5.\nassert api.jvp(square_add_prim, (2., 10.), (1., 1.)) == (14., 5.)\n```\n\nExample:\n```text\ncall square_add_prim(Traced<~float32[]>, Traced<~float32[]>)\n  call multiply_add_prim(Traced<~float32[]>, Traced<~float32[]>, Traced<~float32[]>)\n    call multiply_add_value_and_jvp((2.0, 2.0, 10.0), (1.0, 1.0, 1.0))\n      Primal evaluation:\n      call multiply_add_prim(2.0, 2.0, 10.0)\n        call multiply_add_impl(2.0, 2.0, 10.0)\n        |<- multiply_add_impl = 14.0\n      |<- multiply_add_prim = 14.0\n      Tangent evaluation:\n      call multiply_add_prim(2.0, 1.0, 1.0)\n        call multiply_add_impl(2.0, 1.0, 1.0)\n        |<- multiply_add_impl = 3.0\n      |<- multiply_add_prim = 3.0\n      call multiply_add_prim(1.0, 2.0, 3.0)\n        call multiply_add_impl(1.0, 2.0, 3.0)\n        |<- multiply_add_impl = 5.0\n      |<- multiply_add_prim = 5.0\n    |<- multiply_add_value_and_jvp = (14.0, 5.0)\n  |<- multiply_add_prim = Traced<float32[]>\n|<- square_add_prim = Traced<float32[]>\n```\n\nExample:\n```text\nassert api.jit(lambda arg_values, arg_tangents:\n                   api.jvp(square_add_prim, arg_values, arg_tangents))(\n         (2., 10.), (1., 1.)) == (14., 5.)\n```\n\nExample:\n```text\ncall square_add_prim(Traced<~float32[]>, Traced<~float32[]>)\n  call multiply_add_prim(Traced<~float32[]>, Traced<~float32[]>, Traced<~float32[]>)\n    call multiply_add_value_and_jvp((JitTracer(~float32[]), JitTracer(~float32[]), JitTracer(~float32[])), (JitTracer(~float32[]), JitTracer(~float32[]), JitTracer(~float32[])))\n      Primal evaluation:\n      call multiply_add_prim(JitTracer(~float32[]), JitTracer(~float32[]), JitTracer(~float32[]))\n        call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])\n        |<- multiply_add_abstract_eval = float32[]\n      |<- multiply_add_prim = JitTracer(float32[])\n      Tangent evaluation:\n      call multiply_add_prim(JitTracer(~float32[]), JitTracer(~float32[]), JitTracer(~float32[]))\n        call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])\n        |<- multiply_add_abstract_eval = float32[]\n      |<- multiply_add_prim = JitTracer(float32[])\n      call multiply_add_prim(JitTracer(~float32[]), JitTracer(~float32[]), JitTracer(float32[]))\n        call multiply_add_abstract_eval(~float32[], ~float32[], float32[])\n        |<- multiply_add_abstract_eval = float32[]\n      |<- multiply_add_prim = JitTracer(float32[])\n    |<- multiply_add_value_and_jvp = (JitTracer(float32[]), JitTracer(float32[]))\n  |<- multiply_add_prim = Traced<float32[]>\n|<- square_add_prim = Traced<float32[]>\ncall multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x7a00d8247ed0>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x7a00d80a2f20>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x7a00d80a2f60>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x7a00d8262dc0>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x7a00db66c940>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x7a00d8262ed0>, all_default_mem_kind=True, lowering_cache={}, cached_primitive_lowerings={}, sharding_attr_cache={}, aval_to_ir_types_cache={ShapedArray(float32[], weak_type=True): RankedTensorType(tensor<f32>), ShapedArray(float32[]): RankedTensorType(tensor<f32>), ShapedArray(int32[]): RankedTensorType(tensor<i32>)}, pallas_lowering_cache={}, pallas_collective_id_mapping=CollectiveIdMapping(auto={}, manual={}, all_ids=set()), traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x7a00d80a70b0>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x7a00d8263bb0>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), remove_size_one_mesh_axis=False, xla_metadata=None), platforms=None), BlockArgument(<block argument> of type 'tensor<f32>' at index: 0), BlockArgument(<block argument> of type 'tensor<f32>' at index: 1), BlockArgument(<block argument> of type 'tensor<f32>' at index: 2))\n|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x7a00d80a80b0>]\ncall multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x7a00d8247ed0>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x7a00d80a2f20>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x7a00d80a2f60>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x7a00d8262dc0>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x7a00db66c940>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x7a00d8262ed0>, all_default_mem_kind=True, lowering_cache={LoweringCacheKey(primitive=multiply_add, eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), remove_size_one_mesh_axis=False, xla_metadata=None), avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), effects=frozenset(), params=(), platforms=('cpu',)): LoweringCacheValue(func=<jaxlib.mlir.dialects.func.FuncOp object at 0x7a00d80a3200>, flat_output_types=[RankedTensorType(tensor<f32>)], output_treedef=PyTreeDef([*]), const_args=(), const_arg_avals=(), inline=True)}, cached_primitive_lowerings={}, sharding_attr_cache={}, aval_to_ir_types_cache={ShapedArray(float32[], weak_type=True): RankedTensorType(tensor<f32>), ShapedArray(float32[]): RankedTensorType(tensor<f32>), ShapedArray(int32[]): RankedTensorType(tensor<i32>)}, pallas_lowering_cache={}, pallas_collective_id_mapping=CollectiveIdMapping(auto={}, manual={}, all_ids=set()), traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x7a00d80a70b0>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True), ShapedArray(float32[])), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x7a00d8263ca0>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), remove_size_one_mesh_axis=False, xla_metadata=None), platforms=None), BlockArgument(<block argument> of type 'tensor<f32>' at index: 0), BlockArgument(<block argument> of type 'tensor<f32>' at index: 1), BlockArgument(<block argument> of type 'tensor<f32>' at index: 2))\n|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x7a00d80aa8b0>]\n```\n\nExample:\n```text\n# This is reverse differentiation w.r.t. the first argument of `square_add_prim`\nwith expectNotImplementedError():\n  api.grad(square_add_prim)(2., 10.)\n```\n\nExample:\n```text\ncall square_add_prim(GradTracer(primal=2.0, typeof(tangent)=f32[]), 10.0)\n  call multiply_add_prim(GradTracer(primal=2.0, typeof(tangent)=f32[]), GradTracer(primal=2.0, typeof(tangent)=f32[]), 10.0)\n    call multiply_add_value_and_jvp((2.0, 2.0, 10.0), (Traced<~float32[]>, Traced<~float32[]>, Zero(~float32[])))\n      Primal evaluation:\n      call multiply_add_prim(2.0, 2.0, 10.0)\n        call multiply_add_impl(2.0, 2.0, 10.0)\n        |<- multiply_add_impl = 14.0\n      |<- multiply_add_prim = 14.0\n      Tangent evaluation:\n      call multiply_add_prim(2.0, Traced<~float32[]>, 0.0)\n        call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])\n        |<- multiply_add_abstract_eval = float32[]\n      |<- multiply_add_prim = Traced<float32[]>\n      call multiply_add_prim(Traced<~float32[]>, 2.0, Traced<float32[]>)\n        call multiply_add_abstract_eval(~float32[], ~float32[], float32[])\n        |<- multiply_add_abstract_eval = float32[]\n      |<- multiply_add_prim = Traced<float32[]>\n    |<- multiply_add_value_and_jvp = (14.0, Traced<float32[]>)\n    call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])\n    |<- multiply_add_abstract_eval = float32[]\n    call multiply_add_abstract_eval(~float32[], ~float32[], float32[])\n    |<- multiply_add_abstract_eval = float32[]\n  |<- multiply_add_prim = GradTracer(primal=14.0, typeof(tangent)=f32[])\n|<- square_add_prim = GradTracer(primal=14.0, typeof(tangent)=f32[])\n\nFound expected exception:\n```\n\nExample:\n```text\nTraceback (most recent call last):\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/interpreters/ad.py\", line 292, in get_primitive_transpose\n    return primitive_transposes[p]\n           ~~~~~~~~~~~~~~~~~~~~^^^\nKeyError: multiply_add\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n  File \"<frozen runpy>\", line 198, in _run_module_as_main\n  File \"<frozen runpy>\", line 88, in _run_code\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/ipykernel_launcher.py\", line 18, in <module>\n    app.launch_new_instance()\njax._src.source_info_util.JaxStackTraceBeforeTransformation: NotImplementedError: Transpose rule (for reverse-mode differentiation) for 'multiply_add' not implemented\n\nThe preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.\n\n--------------------\n\nThe above exception was the direct cause of the following exception:\n\nTraceback (most recent call last):\n  File \"/tmp/ipykernel_2054/2155094905.py\", line 3, in <module>\n    api.grad(square_add_prim)(2., 10.)\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/traceback_util.py\", line 194, in reraise_with_filtered_traceback\n    return fun(*args, **kwargs)  # pyrefly: ignore[not-callable]\n           ^^^^^^^^^^^^^^^^^^^^\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/api.py\", line 481, in grad_f\n    _, g = value_and_grad_f(*args, **kwargs)\n           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nNotImplementedError: Transpose rule (for reverse-mode differentiation) for 'multiply_add' not implemented\n```\n\nExample:\n```text\na = xt * 4.\n   b = 2. * yt\n   c = a + b\n   ft = c + yt\n```\n\nExample:\n```text\n# Initialize cotangents of inputs and intermediate variables:\n  xct = yct = act = bct = cct = 0.\n  # Initialize cotangent of the output:\n  fct = 1.\n  # Process `ft = c + yt`:\n  cct += fct\n  yct += fct\n  # Process `c = a + b`:\n  act += cct\n  bct += cct\n  # Process `b = 2. * yt`:\n  yct += 2. * bct\n  # Process `a = xt * 4.`:\n  xct += act * 4.\n```\n\nExample:\n```text\np_transpose(out_ct, x, _, _) = (None, out_ct*cy, out_ct*cz)\n```\n\nExample:\n```text\nadd_transpose(out_ct, _, _) = (out_ct, out_ct)\n mult_transpose(out_ct, x, _) = (None, x * out_ct)\n mult_transpose(out_ct, _, y) = (out_ct * y, None)\n```\n\nExample:\n```text\n@trace(\"multiply_add_transpose\")\ndef multiply_add_transpose(ct, x, y, z):\n  \"\"\"Evaluates the transpose of a linear primitive.\n\n  This method is only used when computing the backward gradient following\n  `value_and_jvp`, and is only needed for primitives that are used in the JVP\n  calculation for some other primitive. You need a transposition for `multiply_add_prim`,\n  because you have used `multiply_add_prim` in the computation of the `output_tangent` in\n  `multiply_add_value_and_jvp`.\n\n  In this case, multiply_add is not a linear primitive. However, it is used linearly\n  w.r.t. tangents in `multiply_add_value_and_jvp`:\n       `output_tangent(xt, yt, zt) = multiply_add_prim(xt, y, multiply_add_prim(x, yt, zt))`.\n\n  Always one of the first two multiplicative arguments is a constant.\n\n  Args:\n      ct: The cotangent of the output of the primitive.\n      x, y, z: The values of the arguments. The arguments that are used linearly\n        get an ad.UndefinedPrimal value. The other arguments get a constant\n        value.\n\n  Returns:\n      A tuple with the cotangent of the inputs, with the value None\n      corresponding to the constant arguments.\n  \"\"\"\n  if not ad.is_undefined_primal(x):\n    # This use of multiply_add is with a constant \"x\".\n    assert ad.is_undefined_primal(y)\n    ct_y = ad.Zero(y.aval) if type(ct) is ad.Zero else multiply_add_prim(x, ct, lax.full_like(x, 0))\n    res = None, ct_y, ct\n  else:\n    # This use of multiply_add is with a constant \"y\".\n    assert ad.is_undefined_primal(x)\n    ct_x = ad.Zero(x.aval) if type(ct) is ad.Zero else multiply_add_prim(ct, y, lax.full_like(y, 0))\n    res = ct_x, None, ct\n  return res\n\nad.primitive_transposes[multiply_add_p] = multiply_add_transpose\n```\n\nExample:\n```text\nassert api.grad(square_add_prim)(2., 10.) == 4.\n```\n\nExample:\n```text\ncall square_add_prim(GradTracer(primal=2.0, typeof(tangent)=f32[]), 10.0)\n  call multiply_add_prim(GradTracer(primal=2.0, typeof(tangent)=f32[]), GradTracer(primal=2.0, typeof(tangent)=f32[]), 10.0)\n    call multiply_add_value_and_jvp((2.0, 2.0, 10.0), (Traced<~float32[]>, Traced<~float32[]>, Zero(~float32[])))\n      Primal evaluation:\n      call multiply_add_prim(2.0, 2.0, 10.0)\n        call multiply_add_impl(2.0, 2.0, 10.0)\n        |<- multiply_add_impl = 14.0\n      |<- multiply_add_prim = 14.0\n      Tangent evaluation:\n      call multiply_add_prim(2.0, Traced<~float32[]>, 0.0)\n        call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])\n        |<- multiply_add_abstract_eval = float32[]\n      |<- multiply_add_prim = Traced<float32[]>\n      call multiply_add_prim(Traced<~float32[]>, 2.0, Traced<float32[]>)\n        call multiply_add_abstract_eval(~float32[], ~float32[], float32[])\n        |<- multiply_add_abstract_eval = float32[]\n      |<- multiply_add_prim = Traced<float32[]>\n    |<- multiply_add_value_and_jvp = (14.0, Traced<float32[]>)\n  |<- multiply_add_prim = GradTracer(primal=14.0, typeof(tangent)=f32[])\n|<- square_add_prim = GradTracer(primal=14.0, typeof(tangent)=f32[])\ncall multiply_add_transpose(1.0, UndefinedPrimal(~float32[]), 2.0, UndefinedPrimal(float32[]))\n  call multiply_add_prim(1.0, 2.0, 0.0)\n    call multiply_add_impl(1.0, 2.0, 0.0)\n    |<- multiply_add_impl = 2.0\n  |<- multiply_add_prim = 2.0\n|<- multiply_add_transpose = (2.0, None, 1.0)\ncall multiply_add_transpose(1.0, 2.0, UndefinedPrimal(~float32[]), 0.0)\n  call multiply_add_prim(2.0, 1.0, 0.0)\n    call multiply_add_impl(2.0, 1.0, 0.0)\n    |<- multiply_add_impl = 2.0\n  |<- multiply_add_prim = 2.0\n|<- multiply_add_transpose = (None, 2.0, 1.0)\n```\n\nExample:\n```text\nassert api.jit(api.grad(square_add_prim))(2., 10.) == 4.\n```\n\nExample:\n```text\ncall square_add_prim(GradTracer(primal=JitTracer(~float32[]), typeof(tangent)=f32[]), JitTracer(~float32[]))\n  call multiply_add_prim(GradTracer(primal=JitTracer(~float32[]), typeof(tangent)=f32[]), GradTracer(primal=JitTracer(~float32[]), typeof(tangent)=f32[]), JitTracer(~float32[]))\n    call multiply_add_value_and_jvp((JitTracer(~float32[]), JitTracer(~float32[]), JitTracer(~float32[])), (Traced<~float32[]>, Traced<~float32[]>, Zero(~float32[])))\n      Primal evaluation:\n      call multiply_add_prim(JitTracer(~float32[]), JitTracer(~float32[]), JitTracer(~float32[]))\n        call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])\n        |<- multiply_add_abstract_eval = float32[]\n      |<- multiply_add_prim = JitTracer(float32[])\n      Tangent evaluation:\n      call multiply_add_prim(JitTracer(~float32[]), Traced<~float32[]>, JitTracer(~float32[]))\n        call multiply_add_abstract_eval(~float32[], ~float32[], ~float32[])\n        |<- multiply_add_abstract_eval = float32[]\n      |<- multiply_add_prim = Traced<float32[]>\n      call multiply_add_prim(Traced<~float32[]>, JitTracer(~float32[]), Traced<float32[]>)\n        call multiply_add_abstract_eval(~float32[], ~float32[], float32[])\n        |<- multiply_add_abstract_eval = float32[]\n      |<- multiply_add_prim = Traced<float32[]>\n    |<- multiply_add_value_and_jvp = (JitTracer(float32[]), Traced<float32[]>)\n  |<- multiply_add_prim = GradTracer(primal=JitTracer(float32[]), typeof(tangent)=f32[])\n|<- square_add_prim = GradTracer(primal=JitTracer(float32[]), typeof(tangent)=f32[])\ncall multiply_add_transpose(JitTracer(float32[]), UndefinedPrimal(~float32[]), JitTracer(~float32[]), UndefinedPrimal(float32[]))\n  call multiply_add_prim(JitTracer(float32[]), JitTracer(~float32[]), JitTracer(~float32[]))\n    call multiply_add_abstract_eval(float32[], ~float32[], ~float32[])\n    |<- multiply_add_abstract_eval = float32[]\n  |<- multiply_add_prim = JitTracer(float32[])\n|<- multiply_add_transpose = (JitTracer(float32[]), None, JitTracer(float32[]))\ncall multiply_add_transpose(JitTracer(float32[]), JitTracer(~float32[]), UndefinedPrimal(~float32[]), JitTracer(~float32[]))\n  call multiply_add_prim(JitTracer(~float32[]), JitTracer(float32[]), JitTracer(~float32[]))\n    call multiply_add_abstract_eval(~float32[], float32[], ~float32[])\n    |<- multiply_add_abstract_eval = float32[]\n  |<- multiply_add_prim = JitTracer(float32[])\n|<- multiply_add_transpose = (None, JitTracer(float32[]), JitTracer(float32[]))\ncall multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x7a00d8247ed0>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x7a00d80eda30>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x7a00d80edca0>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x7a00d80b78d0>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x7a00db66c940>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x7a00d80b7740>, all_default_mem_kind=True, lowering_cache={}, cached_primitive_lowerings={}, sharding_attr_cache={}, aval_to_ir_types_cache={ShapedArray(float32[], weak_type=True): RankedTensorType(tensor<f32>), ShapedArray(float32[]): RankedTensorType(tensor<f32>), ShapedArray(int32[]): RankedTensorType(tensor<i32>)}, pallas_lowering_cache={}, pallas_collective_id_mapping=CollectiveIdMapping(auto={}, manual={}, all_ids=set()), traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x7a00d80e65b0>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[]), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x7a00d80b8130>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), remove_size_one_mesh_axis=False, xla_metadata=None), platforms=None), BlockArgument(<block argument> of type 'tensor<f32>' at index: 0), BlockArgument(<block argument> of type 'tensor<f32>' at index: 1), BlockArgument(<block argument> of type 'tensor<f32>' at index: 2))\n|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x7a00d8303170>]\ncall multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x7a00d8247ed0>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x7a00d80eda30>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x7a00d80edca0>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x7a00d80b78d0>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x7a00db66c940>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x7a00d80b7740>, all_default_mem_kind=True, lowering_cache={LoweringCacheKey(primitive=multiply_add, eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), remove_size_one_mesh_axis=False, xla_metadata=None), avals_in=(ShapedArray(float32[]), ShapedArray(float32[], weak_type=True), ShapedArray(float32[], weak_type=True)), effects=frozenset(), params=(), platforms=('cpu',)): LoweringCacheValue(func=<jaxlib.mlir.dialects.func.FuncOp object at 0x7a00d80ee080>, flat_output_types=[RankedTensorType(tensor<f32>)], output_treedef=PyTreeDef([*]), const_args=(), const_arg_avals=(), inline=True)}, cached_primitive_lowerings={}, sharding_attr_cache={}, aval_to_ir_types_cache={ShapedArray(float32[], weak_type=True): RankedTensorType(tensor<f32>), ShapedArray(float32[]): RankedTensorType(tensor<f32>), ShapedArray(int32[]): RankedTensorType(tensor<i32>)}, pallas_lowering_cache={}, pallas_collective_id_mapping=CollectiveIdMapping(auto={}, manual={}, all_ids=set()), traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x7a00d80e65b0>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[], weak_type=True), ShapedArray(float32[]), ShapedArray(float32[], weak_type=True)), avals_out=[ShapedArray(float32[])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x7a00d80b84f0>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), remove_size_one_mesh_axis=False, xla_metadata=None), platforms=None), BlockArgument(<block argument> of type 'tensor<f32>' at index: 0), BlockArgument(<block argument> of type 'tensor<f32>' at index: 1), BlockArgument(<block argument> of type 'tensor<f32>' at index: 2))\n|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x7a00d80f0ab0>]\n```\n\nExample:\n```text\n# The arguments are two vectors instead of two scalars.\nwith expectNotImplementedError():\n  api.vmap(square_add_prim, in_axes=0, out_axes=0)(np.array([2., 3.]),\n                                               np.array([10., 20.]))\n```\n\nExample:\n```text\ncall square_add_prim(Traced<float32[]>, Traced<float32[]>)\n  call multiply_add_prim(Traced<float32[]>, Traced<float32[]>, Traced<float32[]>)\n\nFound expected exception:\n```\n\nExample:\n```text\nTraceback (most recent call last):\n  File \"/tmp/ipykernel_2054/1080163607.py\", line 3, in <module>\n    api.vmap(square_add_prim, in_axes=0, out_axes=0)(np.array([2., 3.]),\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/traceback_util.py\", line 194, in reraise_with_filtered_traceback\n    return fun(*args, **kwargs)  # pyrefly: ignore[not-callable]\n           ^^^^^^^^^^^^^^^^^^^^\n  File \"/home/docs/checkouts/readthedocs.org/user_builds/jax/envs/latest/lib/python3.12/site-packages/jax/_src/api.py\", line 1239, in vmap_f\n    out_flat, inferred_out_axes = batching.batch(\n                                  ^^^^^^^^^^^^^^^\nNotImplementedError: Batching rule for 'multiply_add' not implemented\n```\n\nExample:\n```text\nfrom jax.interpreters import batching\n\n@trace(\"multiply_add_batch\")\ndef multiply_add_batch(vector_arg_values, batch_axes):\n  \"\"\"Computes the batched version of the primitive.\n\n  This must be a JAX-traceable function.\n\n  Since the `multiply_add primitive` already operates point-wise on arbitrary\n  dimension tensors, to batch it you can use the primitive itself. This works as\n  long as both the inputs have the same dimensions and are batched along the\n  same axes. The result is batched along the axis that the inputs are batched.\n\n  Args:\n    vector_arg_values: A tuple of two arguments, each being a tensor of matching\n      shape.\n    batch_axes: The axes that are being batched. See vmap documentation.\n\n  Returns:\n    A tuple of the result, and the result axis that was batched.\n  \"\"\"\n  assert batch_axes[0] == batch_axes[1]\n  assert batch_axes[0] == batch_axes[2]\n  _trace(\"Using multiply_add to compute the batch:\")\n  res = multiply_add_prim(*vector_arg_values)\n  return res, batch_axes[0]\n\n\nbatching.primitive_batchers[multiply_add_p] = multiply_add_batch\n```\n\nExample:\n```text\nassert np.allclose(api.vmap(square_add_prim, in_axes=0, out_axes=0)(\n  np.array([2., 3.]),\n  np.array([10., 20.])),\n  [14., 29.])\n```\n\nExample:\n```text\ncall square_add_prim(Traced<float32[]>, Traced<float32[]>)\n  call multiply_add_prim(Traced<float32[]>, Traced<float32[]>, Traced<float32[]>)\n    call multiply_add_batch(([2. 3.], [2. 3.], [10. 20.]), (0, 0, 0))\n      Using multiply_add to compute the batch:\n      call multiply_add_prim([2. 3.], [2. 3.], [10. 20.])\n        call multiply_add_impl([2. 3.], [2. 3.], [10. 20.])\n        |<- multiply_add_impl = [14. 29.]\n      |<- multiply_add_prim = [14. 29.]\n    |<- multiply_add_batch = ([14. 29.], 0)\n  |<- multiply_add_prim = Traced<float32[]>\n|<- square_add_prim = Traced<float32[]>\n```\n\nExample:\n```text\nassert np.allclose(api.jit(api.vmap(square_add_prim, in_axes=0, out_axes=0))\n                    (np.array([2., 3.]),\n                     np.array([10., 20.])),\n                    [14., 29.])\n```\n\nExample:\n```text\ncall square_add_prim(Traced<float32[]>, Traced<float32[]>)\n  call multiply_add_prim(Traced<float32[]>, Traced<float32[]>, Traced<float32[]>)\n    call multiply_add_batch((JitTracer(float32[2]), JitTracer(float32[2]), JitTracer(float32[2])), (0, 0, 0))\n      Using multiply_add to compute the batch:\n      call multiply_add_prim(JitTracer(float32[2]), JitTracer(float32[2]), JitTracer(float32[2]))\n        call multiply_add_abstract_eval(float32[2], float32[2], float32[2])\n        |<- multiply_add_abstract_eval = float32[2]\n      |<- multiply_add_prim = JitTracer(float32[2])\n    |<- multiply_add_batch = (JitTracer(float32[2]), 0)\n  |<- multiply_add_prim = Traced<float32[]>\n|<- square_add_prim = Traced<float32[]>\ncall multiply_add_lowering(LoweringRuleContext(module_context=ModuleContext(context=<jax._src.interpreters.mlir.JaxIrContext object at 0x7a00d8247ed0>, module=<jaxlib.mlir._mlir_libs._mlir.ir.Module object at 0x7a00d80ee5c0>, ip=<jaxlib.mlir._mlir_libs._mlir.ir.InsertionPoint object at 0x7a00d80ee600>, symbol_table=<jaxlib.mlir._mlir_libs._mlir.ir.SymbolTable object at 0x7a00d80b8960>, platforms=('cpu',), backend=<jaxlib._jax.Client object at 0x7a00db66c940>, axis_context=ShardingContext(num_devices=1, device_assignment=None, abstract_mesh=None), keepalives=[], channel_iterator=count(2), host_callbacks=[], shape_poly_state=<jax._src.interpreters.mlir.ShapePolyLoweringState object at 0x7a00d80b8a40>, all_default_mem_kind=True, lowering_cache={}, cached_primitive_lowerings={}, sharding_attr_cache={}, aval_to_ir_types_cache={ShapedArray(float32[2]): RankedTensorType(tensor<2xf32>), ShapedArray(int32[]): RankedTensorType(tensor<i32>)}, pallas_lowering_cache={}, pallas_collective_id_mapping=CollectiveIdMapping(auto={}, manual={}, all_ids=set()), traceback_caches=TracebackCaches(traceback_to_location_cache=<jaxlib.mlir._mlir_libs._jax_mlir_ext.TracebackToLocationCache object at 0x7a00d80e55b0>, canonical_name_cache={}), lowering_parameters=LoweringParameters(override_lowering_rules=None, global_constant_computation=False, for_export=False, export_ignore_forward_compatibility=False, hoist_constants_as_args=False)), name_stack=NameStack(stack=()), traceback=None, primitive=multiply_add, avals_in=(ShapedArray(float32[2]), ShapedArray(float32[2]), ShapedArray(float32[2])), avals_out=[ShapedArray(float32[2])], tokens_in=<jax._src.interpreters.mlir.TokenSet object at 0x7a00d80b94b0>, tokens_out=None, const_lowering={}, axis_size_env=None, dim_var_values=[], jaxpr_eqn_ctx=JaxprEqnContext(compute_type=None, threefry_partitionable=True, cur_abstract_mesh=AbstractMesh((), axis_types=()), remove_size_one_mesh_axis=False, xla_metadata=None), platforms=None), BlockArgument(<block argument> of type 'tensor<2xf32>' at index: 0), BlockArgument(<block argument> of type 'tensor<2xf32>' at index: 1), BlockArgument(<block argument> of type 'tensor<2xf32>' at index: 2))\n|<- multiply_add_lowering = [<jaxlib.mlir._mlir_libs._mlir.ir.OpResult object at 0x7a00d80f1af0>]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.835Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":53,"totalLines":877,"estimatedTokens":12314}}116{"id":"doc-generalized_convolutions_in_jax_jax_documentatio-8372aa40","source":"documentation","title":"Generalized convolutions in JAX — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/convolutions.html","text":"Example:\n```text\nimport matplotlib.pyplot as plt\n\nfrom jax import random\nimport jax.numpy as jnp\nimport numpy as np\n\nkey = random.key(1701)\n\nx = jnp.linspace(0, 10, 500)\ny = jnp.sin(x) + 0.2 * random.normal(key, shape=(500,))\n\nwindow = jnp.ones(10) / 10\ny_smooth = jnp.convolve(y, window, mode='same')\n\nplt.plot(x, y, 'lightgray')\nplt.plot(x, y_smooth, 'black');\n```\n\nExample:\n```text\nfrom scipy import datasets\nimport jax.scipy as jsp\n\nfig, ax = plt.subplots(1, 3, figsize=(12, 5))\n\n# Load a sample image; compute mean() to convert from RGB to grayscale.\nimage = jnp.array(datasets.face().mean(-1))\nax[0].imshow(image, cmap='binary_r')\nax[0].set_title('original')\n\n# Create a noisy version by adding random Gaussian noise\nkey = random.key(1701)\nnoisy_image = image + 50 * random.normal(key, image.shape)\nax[1].imshow(noisy_image, cmap='binary_r')\nax[1].set_title('noisy')\n\n# Smooth the noisy image with a 2D Gaussian smoothing kernel.\nx = jnp.linspace(-3, 3, 7)\nwindow = jsp.stats.norm.pdf(x) * jsp.stats.norm.pdf(x[:, None])\nsmooth_image = jsp.signal.convolve(noisy_image, window, mode='same')\nax[2].imshow(smooth_image, cmap='binary_r')\nax[2].set_title('smoothed');\n```\n\nExample:\n```text\n# 2D kernel - HWIO layout\nkernel = jnp.zeros((3, 3, 3, 3), dtype=jnp.float32)\nkernel += jnp.array([[1, 1, 0],\n                     [1, 0,-1],\n                     [0,-1,-1]])[:, :, jnp.newaxis, jnp.newaxis]\n\nprint(\"Edge Conv kernel:\")\nplt.imshow(kernel[:, :, 0, 0]);\n```\n\nExample:\n```text\nEdge Conv kernel:\n```\n\nExample:\n```text\n# NHWC layout\nimg = jnp.zeros((1, 200, 198, 3), dtype=jnp.float32)\nfor k in range(3):\n  x = 30 + 60*k\n  y = 20 + 60*k\n  img = img.at[0, x:x+10, y:y+10, k].set(1.0)\n\nprint(\"Original Image:\")\nplt.imshow(img[0]);\n```\n\nExample:\n```text\nOriginal Image:\n```\n\nExample:\n```text\nfrom jax import lax\nout = lax.conv(jnp.transpose(img,[0,3,1,2]),    # lhs = NCHW image tensor\n               jnp.transpose(kernel,[3,2,0,1]), # rhs = OIHW conv kernel tensor\n               (1, 1),  # window strides\n               'SAME') # padding mode\nprint(\"out shape: \", out.shape)\nprint(\"First output channel:\")\nplt.figure(figsize=(10,10))\nplt.imshow(np.array(out)[0,0,:,:]);\n```\n\nExample:\n```text\nout shape:  (1, 3, 200, 198)\nFirst output channel:\n```\n\nExample:\n```text\nout = lax.conv_with_general_padding(\n  jnp.transpose(img,[0,3,1,2]),    # lhs = NCHW image tensor\n  jnp.transpose(kernel,[3,2,0,1]), # rhs = OIHW conv kernel tensor\n  (1, 1),  # window strides\n  ((2,2),(2,2)), # general padding 2x2\n  (1,1),  # lhs/image dilation\n  (1,1))  # rhs/kernel dilation\nprint(\"out shape: \", out.shape)\nprint(\"First output channel:\")\nplt.figure(figsize=(10,10))\nplt.imshow(np.array(out)[0,0,:,:]);\n```\n\nExample:\n```text\nout shape:  (1, 3, 202, 200)\nFirst output channel:\n```\n\nExample:\n```text\ndn = lax.conv_dimension_numbers(img.shape,     # only ndim matters, not shape\n                                kernel.shape,  # only ndim matters, not shape\n                                ('NHWC', 'HWIO', 'NHWC'))  # the important bit\nprint(dn)\n```\n\nExample:\n```text\nConvDimensionNumbers(lhs_spec=(0, 3, 1, 2), rhs_spec=(3, 2, 0, 1), out_spec=(0, 3, 1, 2))\n```\n\nExample:\n```text\nout = lax.conv_general_dilated(img,    # lhs = image tensor\n                               kernel, # rhs = conv kernel tensor\n                               (1,1),  # window strides\n                               'SAME', # padding mode\n                               (1,1),  # lhs/image dilation\n                               (1,1),  # rhs/kernel dilation\n                               dn)     # dimension_numbers = lhs, rhs, out dimension permutation\nprint(\"out shape: \", out.shape)\nprint(\"First output channel:\")\nplt.figure(figsize=(10,10))\nplt.imshow(np.array(out)[0,:,:,0]);\n```\n\nExample:\n```text\nout shape:  (1, 200, 198, 3)\nFirst output channel:\n```\n\nExample:\n```text\nout = lax.conv_general_dilated(img,     # lhs = image tensor\n                               kernel,  # rhs = conv kernel tensor\n                               (1,1),   # window strides\n                               'VALID', # padding mode\n                               (1,1),   # lhs/image dilation\n                               (1,1),   # rhs/kernel dilation\n                               dn)      # dimension_numbers = lhs, rhs, out dimension permutation\nprint(\"out shape: \", out.shape, \"DIFFERENT from above!\")\nprint(\"First output channel:\")\nplt.figure(figsize=(10,10))\nplt.imshow(np.array(out)[0,:,:,0]);\n```\n\nExample:\n```text\nout shape:  (1, 198, 196, 3) DIFFERENT from above!\nFirst output channel:\n```\n\nExample:\n```text\nout = lax.conv_general_dilated(img,    # lhs = image tensor\n                               kernel, # rhs = conv kernel tensor\n                               (2,2),  # window strides\n                               'SAME', # padding mode\n                               (1,1),  # lhs/image dilation\n                               (1,1),  # rhs/kernel dilation\n                               dn)     # dimension_numbers = lhs, rhs, out dimension permutation\nprint(\"out shape: \", out.shape, \" <-- half the size of above\")\nplt.figure(figsize=(10,10))\nprint(\"First output channel:\")\nplt.imshow(np.array(out)[0,:,:,0]);\n```\n\nExample:\n```text\nout shape:  (1, 100, 99, 3)  <-- half the size of above\nFirst output channel:\n```\n\nExample:\n```text\nout = lax.conv_general_dilated(img,     # lhs = image tensor\n                               kernel,  # rhs = conv kernel tensor\n                               (1,1),   # window strides\n                               'VALID', # padding mode\n                               (1,1),   # lhs/image dilation\n                               (12,12), # rhs/kernel dilation\n                               dn)      # dimension_numbers = lhs, rhs, out dimension permutation\nprint(\"out shape: \", out.shape)\nplt.figure(figsize=(10,10))\nprint(\"First output channel:\")\nplt.imshow(np.array(out)[0,:,:,0]);\n```\n\nExample:\n```text\nout shape:  (1, 176, 174, 3)\nFirst output channel:\n```\n\nExample:\n```text\nout = lax.conv_general_dilated(img,               # lhs = image tensor\n                               kernel,            # rhs = conv kernel tensor\n                               (1,1),             # window strides\n                               ((0, 0), (0, 0)),  # padding mode\n                               (2,2),             # lhs/image dilation\n                               (1,1),             # rhs/kernel dilation\n                               dn)                # dimension_numbers = lhs, rhs, out dimension permutation\nprint(\"out shape: \", out.shape, \"<-- larger than original!\")\nplt.figure(figsize=(10,10))\nprint(\"First output channel:\")\nplt.imshow(np.array(out)[0,:,:,0]);\n```\n\nExample:\n```text\nout shape:  (1, 397, 393, 3) <-- larger than original!\nFirst output channel:\n```\n\nExample:\n```text\n# The following is equivalent to tensorflow:\n# N,H,W,C = img.shape\n# out = tf.nn.conv2d_transpose(img, kernel, (N,2*H,2*W,C), (1,2,2,1))\n\n# transposed conv = 180deg kernel rotation plus LHS dilation\n# rotate kernel 180deg:\nkernel_rot = jnp.rot90(jnp.rot90(kernel, axes=(0,1)), axes=(0,1))\n# need a custom output padding:\npadding = ((2, 1), (2, 1))\nout = lax.conv_general_dilated(img,     # lhs = image tensor\n                               kernel_rot,  # rhs = conv kernel tensor\n                               (1,1),   # window strides\n                               padding, # padding mode\n                               (2,2),   # lhs/image dilation\n                               (1,1),   # rhs/kernel dilation\n                               dn)      # dimension_numbers = lhs, rhs, out dimension permutation\nprint(\"out shape: \", out.shape, \"<-- transposed_conv\")\nplt.figure(figsize=(10,10))\nprint(\"First output channel:\")\nplt.imshow(np.array(out)[0,:,:,0]);\n```\n\nExample:\n```text\nout shape:  (1, 400, 396, 3) <-- transposed_conv\nFirst output channel:\n```\n\nExample:\n```text\n# 1D kernel - WIO layout\nkernel = jnp.array([[[1, 0, -1], [-1,  0,  1]],\n                    [[1, 1,  1], [-1, -1, -1]]],\n                    dtype=jnp.float32).transpose([2,1,0])\n# 1D data - NWC layout\ndata = np.zeros((1, 200, 2), dtype=jnp.float32)\nfor i in range(2):\n  for k in range(2):\n      x = 35*i + 30 + 60*k\n      data[0, x:x+30, k] = 1.0\n\nprint(\"in shapes:\", data.shape, kernel.shape)\n\nplt.figure(figsize=(10,5))\nplt.plot(data[0]);\ndn = lax.conv_dimension_numbers(data.shape, kernel.shape,\n                                ('NWC', 'WIO', 'NWC'))\nprint(dn)\n\nout = lax.conv_general_dilated(data,   # lhs = image tensor\n                               kernel, # rhs = conv kernel tensor\n                               (1,),   # window strides\n                               'SAME', # padding mode\n                               (1,),   # lhs/image dilation\n                               (1,),   # rhs/kernel dilation\n                               dn)     # dimension_numbers = lhs, rhs, out dimension permutation\nprint(\"out shape: \", out.shape)\nplt.figure(figsize=(10,5))\nplt.plot(out[0]);\n```\n\nExample:\n```text\nin shapes: (1, 200, 2) (3, 2, 2)\nConvDimensionNumbers(lhs_spec=(0, 2, 1), rhs_spec=(2, 1, 0), out_spec=(0, 2, 1))\nout shape:  (1, 200, 2)\n```\n\nExample:\n```text\nimport matplotlib as mpl\n\n# Random 3D kernel - HWDIO layout\nkernel = jnp.array([\n  [[0, 0,  0], [0,  1,  0], [0,  0,   0]],\n  [[0, -1, 0], [-1, 0, -1], [0,  -1,  0]],\n  [[0, 0,  0], [0,  1,  0], [0,  0,   0]]],\n  dtype=jnp.float32)[:, :, :, jnp.newaxis, jnp.newaxis]\n\n# 3D data - NHWDC layout\ndata = jnp.zeros((1, 30, 30, 30, 1), dtype=jnp.float32)\nx, y, z = np.mgrid[0:1:30j, 0:1:30j, 0:1:30j]\ndata += (jnp.sin(2*x*jnp.pi)*jnp.cos(2*y*jnp.pi)*jnp.cos(2*z*jnp.pi))[None,:,:,:,None]\n\nprint(\"in shapes:\", data.shape, kernel.shape)\ndn = lax.conv_dimension_numbers(data.shape, kernel.shape,\n                                ('NHWDC', 'HWDIO', 'NHWDC'))\nprint(dn)\n\nout = lax.conv_general_dilated(data,    # lhs = image tensor\n                               kernel,  # rhs = conv kernel tensor\n                               (1,1,1), # window strides\n                               'SAME',  # padding mode\n                               (1,1,1), # lhs/image dilation\n                               (1,1,1), # rhs/kernel dilation\n                               dn)      # dimension_numbers\nprint(\"out shape: \", out.shape)\n\n# Make some simple 3d density plots:\ndef make_alpha(cmap):\n  my_cmap = cmap(jnp.arange(cmap.N))\n  my_cmap[:,-1] = jnp.linspace(0, 1, cmap.N)**3\n  return mpl.colors.ListedColormap(my_cmap)\nmy_cmap = make_alpha(plt.cm.viridis)\nfig = plt.figure()\nax = fig.add_subplot(projection='3d')\nax.scatter(x.ravel(), y.ravel(), z.ravel(), c=data.ravel(), cmap=my_cmap)\nax.axis('off')\nax.set_title('input')\nfig = plt.figure()\nax = fig.add_subplot(projection='3d')\nax.scatter(x.ravel(), y.ravel(), z.ravel(), c=out.ravel(), cmap=my_cmap)\nax.axis('off')\nax.set_title('3D conv output');\n```\n\nExample:\n```text\nin shapes: (1, 30, 30, 30, 1) (3, 3, 3, 1, 1)\nConvDimensionNumbers(lhs_spec=(0, 4, 1, 2, 3), rhs_spec=(4, 3, 0, 1, 2), out_spec=(0, 4, 1, 2, 3))\nout shape:  (1, 30, 30, 30, 1)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.838Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":28,"totalLines":365,"estimatedTokens":2763}}117{"id":"doc-shape_polymorphism_jax_documentation-491d4e88","source":"documentation","title":"Shape polymorphism — JAX documentation","url":"https://docs.jax.dev/en/latest/export/shape_poly.html","text":"Example:\n```text\n>>> import jax\n>>> from jax import export\n>>> from jax import numpy as jnp\n>>> def f(x):  # f: f32[a, b]\n...   return jnp.concatenate([x, x], axis=1)\n\n>>> # We construct symbolic dimension variables.\n>>> a, b = export.symbolic_shape(\"a, b\")\n\n>>> # We can use the symbolic dimensions to construct shapes.\n>>> x_shape = (a, b)\n>>> x_shape\n(a, b)\n\n>>> # Then we export with symbolic shapes:\n>>> exp: export.Exported = export.export(jax.jit(f))(\n...     jax.ShapeDtypeStruct(x_shape, jnp.int32))\n>>> exp.in_avals\n(ShapedArray(int32[a,b]),)\n>>> exp.out_avals\n(ShapedArray(int32[a,2*b]),)\n\n>>> # We can later call with concrete shapes (with a=3 and b=4), without re-tracing `f`.\n>>> res = exp.call(np.ones((3, 4), dtype=np.int32))\n>>> res.shape\n(3, 8)\n```\n\nExample:\n```text\n>>> def f1(x, y): # x: f32[a, 1], y : f32[a, 4]\n...  return x + y\n\n>>> # Assuming you have some actual args with concrete shapes\n>>> x = np.ones((3, 1), dtype=np.int32)\n>>> y = np.ones((3, 4), dtype=np.int32)\n>>> args_specs = export.symbolic_args_specs((x, y), \"a, ...\")\n>>> exp = export.export(jax.jit(f1))(* args_specs)\n>>> exp.in_avals\n(ShapedArray(int32[a,1]), ShapedArray(int32[a,4]))\n```\n\nExample:\n```text\n>>> f = lambda x: jnp.reshape(x, (x.shape[0] * x.shape[1],))\n>>> arg_spec = jax.ShapeDtypeStruct(export.symbolic_shape(\"b, 4\"), jnp.int32)\n>>> exp = export.export(jax.jit(f))(arg_spec)\n>>> exp.out_avals\n(ShapedArray(int32[4*b]),)\n```\n\nExample:\n```text\n>>> exp = export.export(jax.jit(lambda x: jnp.array(x.shape[0]) + x))(\n...     jax.ShapeDtypeStruct(export.symbolic_shape(\"b\"), np.int32))\n>>> exp.call(jnp.arange(3, dtype=np.int32))\nArray([3, 4, 5], dtype=int32)\n\n>>> exp = export.export(jax.jit(lambda x: x.reshape(jnp.array(x.shape[0]) + 2)))(\n...     jax.ShapeDtypeStruct(export.symbolic_shape(\"b\"), np.int32))  \nTraceback (most recent call last):\nTypeError: Shapes must be 1D sequences of concrete values of integer type, got [Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>].\n```\n\nExample:\n```text\n>>> exp = export.export(jax.jit(\n...     lambda x: (5. + x.shape[0],\n...                x.shape[0] - np.arange(5, dtype=jnp.int32),\n...                x + x.shape[0] + jnp.sin(x.shape[0]))))(\n...     jax.ShapeDtypeStruct(export.symbolic_shape(\"b\"), jnp.int32))\n>>> exp.out_avals\n(ShapedArray(float32[], weak_type=True),\n ShapedArray(int32[5]),\n ShapedArray(float32[b], weak_type=True))\n\n>>> exp.call(jnp.ones((3,), jnp.int32))\n (Array(8., dtype=float32, weak_type=True),\n  Array([ 3, 2, 1, 0, -1], dtype=int32),\n  Array([4.14112, 4.14112, 4.14112], dtype=float32, weak_type=True))\n```\n\nExample:\n```text\n>>> exp = export.export(jax.jit(\n...     lambda x: jnp.sum(x, axis=0) / x.shape[0]))(\n...     jax.ShapeDtypeStruct(export.symbolic_shape(\"b, c\"), jnp.int32))\n>>> exp.call(jnp.arange(12, dtype=jnp.int32).reshape((3, 4)))\nArray([4., 5., 6., 7.], dtype=float32)\n```\n\nExample:\n```text\n>>> v, = export.symbolic_shape(\"v,\")\n>>> export.export(jax.jit(lambda x, y: x + y))(      \n...     jax.ShapeDtypeStruct((v,), dtype=np.int32),  \n...     jax.ShapeDtypeStruct((4,), dtype=np.int32))  \nTraceback (most recent call last):\nTypeError: add got incompatible shapes for broadcasting: (v,), (4,).\n\n>>> export.export(jax.jit(lambda x: jnp.matmul(x, x)))(  \n...     jax.ShapeDtypeStruct((v, 4), dtype=np.int32))    \nTraceback (most recent call last):\nTypeError: dot_general requires contracting dimensions to have the same shape, got (4,) and (v,).\n```\n\nExample:\n```text\nimport jax\n>>> export.export(jax.jit(lambda x: 0 if x.shape[0] + 1 >= x.shape[1] else 1))(\n...     jax.ShapeDtypeStruct(export.symbolic_shape(\"a, b\"), dtype=np.int32))  # doctest: +IGNORE_EXCEPTION_DETAIL\nTraceback (most recent call last):\njax._src.export.shape_poly.InconclusiveDimensionOperation: Symbolic dimension comparison 'a + 1' >= 'b' is inconclusive.\nThis error arises for comparison operations with shapes that\nare non-constant, and the result of the operation cannot be represented as\na boolean value for all values of the symbolic dimensions involved.\n```\n\nExample:\n```text\n>>> _ = export.export(jax.jit(lambda x: x[0:16]))(\n...    jax.ShapeDtypeStruct(export.symbolic_shape(\"b + 15\"), dtype=np.int32))\n```\n\nExample:\n```text\n>>> # Introduce dimension variable with constraints.\n>>> a, b = export.symbolic_shape(\"a, b\",\n...                              constraints=(\"a >= b\", \"b >= 16\"))\n>>> _ = export.export(jax.jit(lambda x: x[:x.shape[1], :16]))(\n...    jax.ShapeDtypeStruct((a, b), dtype=np.int32))\n```\n\nExample:\n```text\n>>> # Introduce dimension variable with equality constraints.\n>>> a, b, c, d = export.symbolic_shape(\"a, b, c, d\",\n...                                    constraints=(\"a * b == c + d\",))\n>>> 2 * b * a\n2*d + 2*c\n\n>>> a * b * b\nb*d + b*c\n```\n\nExample:\n```text\nfrom jax import lax\n>>> b, = export.symbolic_shape(\"b\")\n>>> f = lambda x: lax.slice_in_dim(x, 0, x.shape[0] % 3)\n>>> export.export(jax.jit(f))(\n...     jax.ShapeDtypeStruct((b,), dtype=np.int32))  # doctest: +IGNORE_EXCEPTION_DETAIL\nTraceback (most recent call last):\njax._src.export.shape_poly.InconclusiveDimensionOperation: Symbolic dimension comparison 'b' >= 'mod(b, 3)' is inconclusive.\nThis error arises for comparison operations with shapes that\nare non-constant, and the result of the operation cannot be represented as\na boolean value for all values of the symbolic dimensions involved.\n```\n\nExample:\n```text\n>>> b, = export.symbolic_shape(\"b\",\n...                            constraints=[\"b >= mod(b, 3)\"])\n>>> f = lambda x: lax.slice_in_dim(x, 0, x.shape[0] % 3)\n>>> _ = export.export(jax.jit(f))(\n...     jax.ShapeDtypeStruct((b,), dtype=np.int32))\n```\n\nExample:\n```text\n>>> a1, = export.symbolic_shape(\"a,\")\n>>> a2, = export.symbolic_shape(\"a,\", constraints=(\"a >= 8\",))\n\n>>> a1 + a2  \nTraceback (most recent call last):\nValueError: Invalid mixing of symbolic scopes for linear combination.\nExpected  scope 4776451856 created at <doctest shape_poly.md[31]>:1:6 (<module>)\nand found for 'a' (unknown) scope 4776979920 created at <doctest shape_poly.md[32]>:1:6 (<module>) with constraints:\n  a >= 8\n```\n\nExample:\n```text\n>>> a, = export.symbolic_shape(\"a,\", constraints=(\"a >= 8\",))\n>>> b, = export.symbolic_shape(\"b,\", scope=a.scope)  # Reuse the scope of `a`\n\n>>> a + b  # Allowed\nb + a\n```\n\nExample:\n```text\n>>> my_scope = export.SymbolicScope()\n>>> c, = export.symbolic_shape(\"c\", scope=my_scope)\n>>> d, = export.symbolic_shape(\"d\", scope=my_scope)\n>>> c + d  # Allowed\nd + c\n```\n\nExample:\n```text\n>>> def my_top_k(k, x):  # x: i32[4, 10], k <= 10\n...   return lax.top_k(x, k)[0]  # : i32[4, 3]\n>>> x = np.arange(40, dtype=np.int32).reshape((4, 10))\n\n>>> # Export with static `k=3`. Since `k` appears in shapes it must be in `static_argnums`.\n>>> exp_static_k = export.export(jax.jit(my_top_k, static_argnums=0))(3, x)\n>>> exp_static_k.in_avals[0]\nShapedArray(int32[4,10])\n\n>>> exp_static_k.out_avals[0]\nShapedArray(int32[4,3])\n\n>>> # When calling the exported function we pass only the non-static arguments\n>>> exp_static_k.call(x)\nArray([[ 9,  8,  7],\n       [19, 18, 17],\n       [29, 28, 27],\n       [39, 38, 37]], dtype=int32)\n\n>>> # Now attempt to export with symbolic `k` so that we choose `k` after export.\n>>> k, = export.symbolic_shape(\"k\", constraints=[\"k <= 10\"])\n>>> export.export(jax.jit(my_top_k, static_argnums=0))(k, x)  \nTraceback (most recent call last):\nUnexpectedDimVar: \"Encountered dimension variable 'k' that is not appearing in the shapes of the function arguments\n```\n\nExample:\n```text\n>>> def my_top_k_with_dimensions(dimensions, x):  # dimensions: i32[0, k], x: i32[4, 10]\n...   return my_top_k(dimensions.shape[1], x)\n>>> exp = export.export(jax.jit(my_top_k_with_dimensions))(\n...     jax.ShapeDtypeStruct((0, k), dtype=np.int32),\n...     x)\n>>> exp.in_avals\n(ShapedArray(int32[0,k]), ShapedArray(int32[4,10]))\n\n>>> exp.out_avals[0]\nShapedArray(int32[4,k])\n\n>>> # When we invoke `exp` we must construct and pass an array of shape (0, k)\n>>> exp.call(np.zeros((0, 3), dtype=np.int32), x)\nArray([[ 9,  8,  7],\n       [19, 18, 17],\n       [29, 28, 27],\n       [39, 38, 37]], dtype=int32)\n```\n\nExample:\n```text\n>>> a, = export.symbolic_shape(\"a\")\n>>> export.export(jax.jit(lambda x: x.shape[0]))(\n...    jax.ShapeDtypeStruct((a * a,), dtype=np.int32))  \nTraceback (most recent call last):\nValueError: Cannot solve for values of dimension variables {'a'}.\nWe can only solve linear uni-variate constraints.\nUsing the following polymorphic shapes specifications: args[0].shape = (a^2,).\nUnprocessed specifications: 'a^2' for dimension size args[0].shape[0].\n```\n\nExample:\n```text\n>>> def f(x):  # x: f32[b, b, 2*d]\n...   return x\n>>> exp = export.export(jax.jit(f))(\n...     jax.ShapeDtypeStruct(export.symbolic_shape(\"b, b, 2*d\"), dtype=np.int32))   \n>>> exp.call(np.ones((3, 3, 5), dtype=np.int32))  \nTraceback (most recent call last):\nValueError: Input shapes do not match the polymorphic shapes specification.\nDivision had remainder 1 when computing the value of 'd'.\nUsing the following polymorphic shapes specifications:\n  args[0].shape = (b, b, 2*d).\nObtained dimension variables: 'b' = 3 from specification 'b' for dimension args[0].shape[0] (= 3), .\nPlease see https://docs.jax.dev/en/latest/export/shape_poly.html#shape-assertion-errors for more details.\n```\n\nExample:\n```text\n# Log from python\nJAX_DUMP_IR_TO=/tmp/export.dumps/ TF_CPP_VMODULE=refine_polymorphic_shapes=3 python tests/shape_poly_test.py ShapePolyTest.test_simple_unary -v=3\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.840Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":21,"totalLines":285,"estimatedTokens":2374}}118{"id":"doc-gradient_checkpointing_with_jax_checkpoint_jax_r-305e533d","source":"documentation","title":"Gradient checkpointing with jax.checkpoint (jax.remat) — JAX documentation","url":"https://docs.jax.dev/en/latest/gradient-checkpointing.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\n\ndef g(W, x):\n  y = jnp.dot(W, x)\n  return jnp.sin(y)\n\ndef f(W1, W2, W3, x):\n  x = g(W1, x)\n  x = g(W2, x)\n  x = g(W3, x)\n  return x\n\nW1 = jnp.ones((5, 4))\nW2 = jnp.ones((6, 5))\nW3 = jnp.ones((7, 6))\nx = jnp.ones(4)\n\n# Inspect the 'residual' values to be saved on the forward pass\n# if you were to evaluate `jax.grad(f)(W1, W2, W3, x)`\nfrom jax.ad_checkpoint import print_saved_residuals\nprint_saved_residuals(f, W1, W2, W3, x)\n```\n\nExample:\n```text\nf32[5,4] from the argument W1\nf32[6,5] from the argument W2\nf32[7,6] from the argument W3\nf32[4] from the argument x\nf32[5] output of sin from /tmp/ipykernel_1887/1857807639.py:6:9 (g)\nf32[5] output of cos from /tmp/ipykernel_1887/1857807639.py:6:9 (g)\nf32[6] output of sin from /tmp/ipykernel_1887/1857807639.py:6:9 (g)\nf32[6] output of cos from /tmp/ipykernel_1887/1857807639.py:6:9 (g)\nf32[7] output of cos from /tmp/ipykernel_1887/1857807639.py:6:9 (g)\n```\n\nExample:\n```text\ndef f2(W1, W2, W3, x):\n  x = jax.checkpoint(g)(W1, x)\n  x = jax.checkpoint(g)(W2, x)\n  x = jax.checkpoint(g)(W3, x)\n  return x\n\nprint_saved_residuals(f2, W1, W2, W3, x)\n```\n\nExample:\n```text\nf32[5,4] from the argument W1\nf32[6,5] from the argument W2\nf32[7,6] from the argument W3\nf32[4] from the argument x\nf32[5] output of sin from /tmp/ipykernel_1887/1857807639.py:6:9 (g)\nf32[6] output of sin from /tmp/ipykernel_1887/1857807639.py:6:9 (g)\n```\n\nExample:\n```text\nf3 = jax.checkpoint(f, policy=jax.checkpoint_policies.dots_with_no_batch_dims_saveable)\nprint_saved_residuals(f3, W1, W2, W3, x)\n```\n\nExample:\n```text\nf32[5,4] from the argument W1\nf32[6,5] from the argument W2\nf32[7,6] from the argument W3\nf32[4] from the argument x\nf32[5] output of dot_general from /tmp/ipykernel_1887/1857807639.py:5:6 (g)\nf32[6] output of dot_general from /tmp/ipykernel_1887/1857807639.py:5:6 (g)\nf32[7] output of dot_general from /tmp/ipykernel_1887/1857807639.py:5:6 (g)\n```\n\nExample:\n```text\nfrom jax.ad_checkpoint import checkpoint_name\n\ndef f4(W1, W2, W3, x):\n  x = checkpoint_name(g(W1, x), name='a')\n  x = checkpoint_name(g(W2, x), name='b')\n  x = checkpoint_name(g(W3, x), name='c')\n  return x\n\nf4 = jax.checkpoint(f4, policy=jax.checkpoint_policies.save_only_these_names('a'))\nprint_saved_residuals(f4, W1, W2, W3, x)\n```\n\nExample:\n```text\nf32[5,4] from the argument W1\nf32[6,5] from the argument W2\nf32[7,6] from the argument W3\nf32[4] from the argument x\nf32[5] named 'a' from /tmp/ipykernel_1887/3722338705.py:4:6 (f4)\n```\n\nExample:\n```text\nfrom jax.tree_util import tree_flatten, tree_unflatten\n\nfrom rich.console import Console\nfrom rich.table import Table\nimport rich.text\n\ndef print_fwd_bwd(f, *args, **kwargs) -> None:\n  args, in_tree = tree_flatten((args, kwargs))\n\n  def f_(*args):\n    args, kwargs = tree_unflatten(in_tree, args)\n    return f(*args, **kwargs)\n\n  fwd = jax.make_jaxpr(lambda *args: jax.vjp(f_, *args))(*args).jaxpr\n\n  y, f_vjp = jax.vjp(f_, *args)\n  res, in_tree = tree_flatten(f_vjp)\n\n  def g_(*args):\n    *res, y = args\n    f_vjp = tree_unflatten(in_tree, res)\n    return f_vjp(y)\n\n  bwd = jax.make_jaxpr(g_)(*res, y).jaxpr\n\n  table = Table(show_header=False, show_lines=True, padding=(1, 2, 0, 2), box=None)\n  table.add_row(\"[bold green]forward computation:\",\n                \"[bold green]backward computation:\")\n  table.add_row(rich.text.Text.from_ansi(str(fwd)),\n                rich.text.Text.from_ansi(str(bwd)))\n  console = Console(width=240, force_jupyter=True)\n  console.print(table)\n\ndef _renderable_repr(self):\n  return self.html\nrich.jupyter.JupyterRenderable._repr_html_ = _renderable_repr\n```\n\nExample:\n```text\n# Without using `jax.checkpoint`:\nprint_fwd_bwd(f, W1, W2, W3, x)\n```\n\nExample:\n```text\nforward computation:                                         backward computation:                                                                     \n                                                                                                                                                         \n  { lambda ; a:f32[5,4] b:f32[6,5] c:f32[7,6] d:f32[4]. let    { lambda ; a:f32[5,4] b:f32[6,5] c:f32[7,6] d:f32[4] e:f32[5] f:f32[5] g:f32[6] h:f32[6]  \n      e:f32[5] = dot_general[                                      i:f32[7] j:f32[7]. let                                                                \n        dimension_numbers=(([1], [0]), ([], []))                   k:f32[7] = mul j i                                                                    \n        preferred_element_type=float32                             l:f32[6] = dot_general[                                                               \n      ] a d                                                          dimension_numbers=(([0], [0]), ([], []))                                            \n      f:f32[5] = sin e                                               preferred_element_type=float32                                                      \n      g:f32[5] = cos e                                             ] k c                                                                                 \n      h:f32[6] = dot_general[                                      m:f32[7,6] = dot_general[                                                             \n        dimension_numbers=(([1], [0]), ([], []))                     dimension_numbers=(([], []), ([], []))                                              \n        preferred_element_type=float32                               preferred_element_type=float32                                                      \n      ] b f                                                        ] k h                                                                                 \n      i:f32[6] = sin h                                             n:f32[6] = mul l g                                                                    \n      j:f32[6] = cos h                                             o:f32[5] = dot_general[                                                               \n      k:f32[7] = dot_general[                                        dimension_numbers=(([0], [0]), ([], []))                                            \n        dimension_numbers=(([1], [0]), ([], []))                     preferred_element_type=float32                                                      \n        preferred_element_type=float32                             ] n b                                                                                 \n      ] c i                                                        p:f32[6,5] = dot_general[                                                             \n      l:f32[7] = sin k                                               dimension_numbers=(([], []), ([], []))                                              \n      m:f32[7] = cos k                                               preferred_element_type=float32                                                      \n    in (l, a, b, c, d, g, f, j, i, m) }                            ] n f                                                                                 \n                                                                   q:f32[5] = mul o e                                                                    \n                                                                   r:f32[4] = dot_general[                                                               \n                                                                     dimension_numbers=(([0], [0]), ([], []))                                            \n                                                                     preferred_element_type=float32                                                      \n                                                                   ] q a                                                                                 \n                                                                   s:f32[5,4] = dot_general[                                                             \n                                                                     dimension_numbers=(([], []), ([], []))                                              \n                                                                     preferred_element_type=float32                                                      \n                                                                   ] q d                                                                                 \n                                                                 in (s, p, m, r) }\n```\n\nExample:\n```text\n# Using `jax.checkpoint` with policy=jax.checkpoint_policies.dots_with_no_batch_dims_saveable:\nprint_fwd_bwd(f3, W1, W2, W3, x)\n```\n\nExample:\n```text\nforward computation:                                                   backward computation:                                                               \n                                                                                                                                                             \n  { lambda ; a:f32[5,4] b:f32[6,5] c:f32[7,6] d:f32[4]. let              let jaxpr = { lambda ; a:f32[5] b:f32[6] c:f32[7] d:f32[5,4] e:f32[6,5] f:f32[7,6]  \n      e:f32[5] = dot_general[                                                g:f32[4] h:f32[7]. let                                                          \n        dimension_numbers=(([1], [0]), ([], []))                             i:f32[5] = sin a                                                                \n        preferred_element_type=float32                                       j:f32[5] = cos a                                                                \n      ] a d                                                                  k:f32[6] = sin b                                                                \n      f:f32[5] = reduce_precision[exponent_bits=8 mantissa_bits=23] e        l:f32[6] = cos b                                                                \n      g:f32[5] = sin f                                                       m:f32[7] = cos c                                                                \n      h:f32[6] = dot_general[                                                n:f32[7] = mul h m                                                              \n        dimension_numbers=(([1], [0]), ([], []))                             o:f32[6] = dot_general[                                                         \n        preferred_element_type=float32                                         dimension_numbers=(([0], [0]), ([], []))                                      \n      ] b g                                                                    preferred_element_type=float32                                                \n      i:f32[6] = reduce_precision[exponent_bits=8 mantissa_bits=23] h        ] n f                                                                           \n      j:f32[6] = sin i                                                       p:f32[7,6] = dot_general[                                                       \n      k:f32[7] = dot_general[                                                  dimension_numbers=(([], []), ([], []))                                        \n        dimension_numbers=(([1], [0]), ([], []))                               preferred_element_type=float32                                                \n        preferred_element_type=float32                                       ] n k                                                                           \n      ] c j                                                                  q:f32[6] = mul o l                                                              \n      l:f32[7] = reduce_precision[exponent_bits=8 mantissa_bits=23] k        r:f32[5] = dot_general[                                                         \n      m:f32[7] = sin l                                                         dimension_numbers=(([0], [0]), ([], []))                                      \n    in (m, a, b, c, d, f, i, l) }                                              preferred_element_type=float32                                                \n                                                                             ] q e                                                                           \n                                                                             s:f32[6,5] = dot_general[                                                       \n                                                                               dimension_numbers=(([], []), ([], []))                                        \n                                                                               preferred_element_type=float32                                                \n                                                                             ] q i                                                                           \n                                                                             t:f32[5] = mul r j                                                              \n                                                                             u:f32[4] = dot_general[                                                         \n                                                                               dimension_numbers=(([0], [0]), ([], []))                                      \n                                                                               preferred_element_type=float32                                                \n                                                                             ] t d                                                                           \n                                                                             v:f32[5,4] = dot_general[                                                       \n                                                                               dimension_numbers=(([], []), ([], []))                                        \n                                                                               preferred_element_type=float32                                                \n                                                                             ] t g                                                                           \n                                                                           in (v, s, p, u) } in                                                              \n                                                                         { lambda ; w:f32[5,4] x:f32[6,5] y:f32[7,6] z:f32[4] ba:f32[5] bb:f32[6] bc:f32[7]  \n                                                                             bd:f32[7]. let                                                                  \n                                                                             be:f32[5,4] bf:f32[6,5] bg:f32[7,6] bh:f32[4] = remat2[                         \n                                                                               differentiated=True                                                           \n                                                                               jaxpr=jaxpr                                                                   \n                                                                               policy=DotsSaveable(only_if_no_batch_dims=True)                               \n                                                                               prevent_cse=True                                                              \n                                                                             ] ba bb bc w x y z bd                                                           \n                                                                           in (be, bf, bg, bh) }\n```\n\nExample:\n```text\ndef sin_vjp(x):\n  y = jnp.sin(x)\n  cos_x = jnp.cos(x)\n  return y, lambda y_bar: cos_x * y_bar\n```\n\nExample:\n```text\ndef sin_vjp2(x):\n  y = jnp.sin(x)\n  return y, lambda y_bar: jnp.cos(x) * y_bar\n```\n\nExample:\n```text\ndef f(x):\n  y = g(x)\n  z = h(y)\n  return z\n\ndef f_vjp(x):\n  y, g_vjp = jax.vjp(g, x)\n  z, h_vjp = jax.vjp(h, y)\n  def f_bwd(z_bar):\n    y_bar, = h_vjp(z_bar)\n    x_bar, = g_vjp(y_bar)\n    return x_bar\n  return z, f_bwd\n```\n\nExample:\n```text\ndef f_vjp_checkpoint(x):\n  y = g(x)\n  z, h_vjp = jax.vjp(h, y)\n  def f_bwd2(z_bar):\n    y_bar, = h_vjp(z_bar)\n    _, g_vjp = jax.vjp(g, x)\n    x_bar, = g_vjp(y_bar)\n    return x_bar\n  return z, f_bwd2\n```\n\nExample:\n```text\ndef f_checkpoint(x):\n  y = jax.checkpoint(g)(x)\n  z = h(y)\n  return z\n```\n\nExample:\n```text\ndef f_checkpoint_grad(x):\n  y = g(x)                  # step 1\n  _, h_vjp = jax.vjp(h)(y)  # step 2\n  y_bar, = h_vjp(1.0)       # step 3\n  _, g_vjp = jax.vjp(g, x)  # step 4\n  x_bar, = g_vjp(y_bar)     # step 5\n  return x_bar\n```\n\nExample:\n```text\ndef f_grad_bad(x):\n  _ = f(x)                  # step 1\n  _, f_vjp = jax.vjp(f, x)  # step 2\n  x_bar, = f_vjp(1.0)       # step 3\n  return x_bar\n```\n\nExample:\n```text\ndef f_grad_bad2(x):\n  y, g_vjp = jax.vjp(g, x)  # step 1\n  z = h(y)                  # step 2\n  _, h_vjp = jax.vjp(h, y)  # step 3\n  y_bar, = h_vjp(1.0)       # step 3\n  x_bar, = g_vjp(y_bar)     # step 5\n  return x_bar\n```\n\nExample:\n```text\ndef loss(params, x, y):\n  return jnp.sum((predict(params, x) - y)**2)\n\ndef predict(params, x):\n  *Ws, Wlast = params\n  for W in Ws:\n    x = layer(W, x)\n  x = jnp.dot(Wlast, x)\n  return x\n\ndef layer(W, x):\n  return jnp.sin(jnp.dot(W, x))\n```\n\nExample:\n```text\nW1 = W2 = W3 = jnp.ones((4, 4))\nparams = [W1, W2, W3]\nx = jnp.ones(4)\ny = jnp.ones(4)\n```\n\nExample:\n```text\nprint_saved_residuals(loss, params, x, y)\n```\n\nExample:\n```text\nf32[4,4] from the argument params[0]\nf32[4,4] from the argument params[1]\nf32[4,4] from the argument params[2]\nf32[4] from the argument x\nf32[4] output of sin from /tmp/ipykernel_1887/4230705069.py:12:9 (layer)\nf32[4] output of cos from /tmp/ipykernel_1887/4230705069.py:12:9 (layer)\nf32[4] output of sin from /tmp/ipykernel_1887/4230705069.py:12:9 (layer)\nf32[4] output of cos from /tmp/ipykernel_1887/4230705069.py:12:9 (layer)\nf32[4] output of mul from /tmp/ipykernel_1887/4230705069.py:2:17 (loss)\n```\n\nExample:\n```text\nloss_checkpoint = jax.checkpoint(loss, policy=jax.checkpoint_policies.dots_with_no_batch_dims_saveable)\nprint_saved_residuals(loss_checkpoint, params, x, y)\n```\n\nExample:\n```text\nf32[4,4] from the argument params[0]\nf32[4,4] from the argument params[1]\nf32[4,4] from the argument params[2]\nf32[4] from the argument x\nf32[4] from the argument y\nf32[4] output of dot_general from /tmp/ipykernel_1887/4230705069.py:12:17 (layer)\nf32[4] output of dot_general from /tmp/ipykernel_1887/4230705069.py:12:17 (layer)\nf32[4] output of dot_general from /tmp/ipykernel_1887/4230705069.py:8:6 (predict)\n```\n\nExample:\n```text\nfrom jax.ad_checkpoint import checkpoint_name\n\ndef predict(params, x):\n  *Ws, Wlast = params\n  for i, W in enumerate(Ws):\n    x = layer(W, x)\n    x = checkpoint_name(x, name=f'layer{i}_output')\n  x = jnp.dot(Wlast, x)\n  return x\n```\n\nExample:\n```text\nf32[4,4] from the argument params[0]\nf32[4,4] from the argument params[1]\nf32[4,4] from the argument params[2]\nf32[4] from the argument x\nf32[4] output of cos from /tmp/ipykernel_1887/4230705069.py:12:9 (layer)\nf32[4] named 'layer0_output' from /tmp/ipykernel_1887/178264713.py:7:8 (predict)\nf32[4] output of cos from /tmp/ipykernel_1887/4230705069.py:12:9 (layer)\nf32[4] named 'layer1_output' from /tmp/ipykernel_1887/178264713.py:7:8 (predict)\nf32[4] output of mul from /tmp/ipykernel_1887/4230705069.py:2:17 (loss)\n```\n\nExample:\n```text\nloss_checkpoint2 = jax.checkpoint(loss, policy=jax.checkpoint_policies.save_any_names_but_these('layer1_output'))\nprint_saved_residuals(loss_checkpoint2, params, x, y)\n```\n\nExample:\n```text\nf32[4,4] from the argument params[0]\nf32[4,4] from the argument params[1]\nf32[4,4] from the argument params[2]\nf32[4] from the argument x\nf32[4] from the argument y\n```\n\nExample:\n```text\nfrom jax import checkpoint\n\ndef checkpoint_offload_dot_with_no_batch_dims(self):\n  policy = jax.checkpoint_policies.offload_dot_with_no_batch_dims(\n      \"device\", \"pinned_host\")\n\n  @functools.partial(checkpoint, policy=policy)\n  def f(x):\n    x = jnp.einsum('ij,jk->ik', x, x, precision=lax.Precision.HIGHEST)\n    x = jnp.sin(x)\n    x = jnp.einsum('ij,jk->ik', x, x, precision=lax.Precision.HIGHEST)\n    x = jnp.sin(x)\n    x = jnp.einsum('ij,jk->ik', x, x, precision=lax.Precision.HIGHEST)\n    x = jnp.sin(x)\n    x = jnp.sum(x)\n    return x\n```\n\nExample:\n```text\nfrom jax import checkpoint\nfrom jax.ad_checkpoint import checkpoint_name\nfrom jax._src import test_util as jtu\n\ndef checkpoint_names_saved_offloaded_recomputed(self):\n  mesh = jtu.create_mesh((2,), (\"x\",))\n  shape = (256, 128)\n  np_inp = np.arange(math.prod(shape), dtype=np.float32).reshape(shape)\n  s = NamedSharding(mesh, P(\"x\"))\n  inp = jax.device_put(np_inp, s)\n\n  policy = jax.checkpoint_policies.save_and_offload_only_these_names(\n      names_which_can_be_saved=[\"y\"], names_which_can_be_offloaded=[\"z\"],\n      offload_src='device', offload_dst='pinned_host')\n\n  @functools.partial(checkpoint, policy=policy)\n  def f(x):\n    def g(ys, _):\n      y, _ = ys\n      y = checkpoint_name(jnp.sin(y), \"y\")\n      z = checkpoint_name(jnp.sin(y), \"z\")\n      z = z.T\n      w = checkpoint_name(jnp.sin(z), \"w\")\n      return (w.T, jnp.sum(w)), None\n    _, scan_out = jax.lax.scan(g, (x, np.array(1, dtype=np.float32)), [np_inp])[0]\n    return scan_out\n```\n\nExample:\n```text\ndef chain_compose(funs):\n  def f(x):\n    for fun in funs:\n      x = fun(x)\n    return x\n  return f\n\nf = chain_compose([jnp.sin] * 8)\nprint_saved_residuals(f, 3.)\n```\n\nExample:\n```text\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\n```\n\nExample:\n```text\nf = chain_compose([jnp.sin] * 16)\nprint_saved_residuals(f, 3.)\n```\n\nExample:\n```text\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\nf32[] output of cos from /tmp/ipykernel_1887/410288286.py:4:10 (chain_compose.<locals>.f)\n```\n\nExample:\n```text\ndef recursive_checkpoint(funs):\n  if len(funs) == 1:\n    return funs[0]\n  elif len(funs) == 2:\n    f1, f2 = funs\n    return lambda x: f1(f2(x))\n  else:\n    f1 = recursive_checkpoint(funs[:len(funs)//2])\n    f2 = recursive_checkpoint(funs[len(funs)//2:])\n    return lambda x: f1(jax.checkpoint(f2)(x))\n```\n\nExample:\n```text\nf = recursive_checkpoint([jnp.sin] * 8)\nprint_saved_residuals(f, 3.)\n```\n\nExample:\n```text\nf32[] from the argument x\nf32[] output of sin from /tmp/ipykernel_1887/1943107544.py:6:21 (recursive_checkpoint.<locals>.<lambda>)\nf32[] output of cos from /tmp/ipykernel_1887/1943107544.py:6:24 (recursive_checkpoint.<locals>.<lambda>)\nf32[] output of cos from /tmp/ipykernel_1887/1943107544.py:6:21 (recursive_checkpoint.<locals>.<lambda>)\n```\n\nExample:\n```text\nf = recursive_checkpoint([jnp.sin] * 16)\nprint_saved_residuals(f, 3.)\n```\n\nExample:\n```text\nf32[] from the argument x\nf32[] output of sin from /tmp/ipykernel_1887/1943107544.py:6:21 (recursive_checkpoint.<locals>.<lambda>)\nf32[] output of sin from /tmp/ipykernel_1887/1943107544.py:6:21 (recursive_checkpoint.<locals>.<lambda>)\nf32[] output of cos from /tmp/ipykernel_1887/1943107544.py:6:24 (recursive_checkpoint.<locals>.<lambda>)\nf32[] output of cos from /tmp/ipykernel_1887/1943107544.py:6:21 (recursive_checkpoint.<locals>.<lambda>)\n```\n\nExample:\n```text\nf = chain_compose([jnp.sin] * 8)\nprint_fwd_bwd(f, 3.)\n```\n\nExample:\n```text\nforward computation:                  backward computation:                                                                    \n                                                                                                                                 \n  { lambda ; a:f32[]. let               { lambda ; a:f32[] b:f32[] c:f32[] d:f32[] e:f32[] f:f32[] g:f32[] h:f32[] i:f32[]. let  \n      b:f32[] = sin a                       j:f32[] = mul i h                                                                    \n      c:f32[] = cos a                       k:f32[] = mul j g                                                                    \n      d:f32[] = sin b                       l:f32[] = mul k f                                                                    \n      e:f32[] = cos b                       m:f32[] = mul l e                                                                    \n      f:f32[] = sin d                       n:f32[] = mul m d                                                                    \n      g:f32[] = cos d                       o:f32[] = mul n c                                                                    \n      h:f32[] = sin f                       p:f32[] = mul o b                                                                    \n      i:f32[] = cos f                       q:f32[] = mul p a                                                                    \n      j:f32[] = sin h                     in (q,) }                                                                              \n      k:f32[] = cos h                                                                                                            \n      l:f32[] = sin j                                                                                                            \n      m:f32[] = cos j                                                                                                            \n      n:f32[] = sin l                                                                                                            \n      o:f32[] = cos l                                                                                                            \n      p:f32[] = sin n                                                                                                            \n      q:f32[] = cos n                                                                                                            \n    in (p, c, e, g, i, k, m, o, q) }\n```\n\nExample:\n```text\nf = recursive_checkpoint([jnp.sin] * 8)\nprint_fwd_bwd(f, 3.)\n```\n\nExample:\n```text\nforward computation:                                                              backward computation:                               \n                                                                                                                                        \n  { lambda ; a:f32[]. let                                                           { lambda ; a:f32[] b:f32[] c:f32[] d:f32[]. let     \n      b:f32[] = remat2[                                                                 e:f32[] = mul d c                               \n        differentiated=False                                                            f:f32[] = mul e b                               \n        jaxpr={ lambda ; c:f32[]. let d:f32[] = sin c; e:f32[] = sin d in (e,) }        g:f32[] = remat2[                               \n        policy=None                                                                       differentiated=True                           \n        prevent_cse=True                                                                  jaxpr={ lambda ; h:f32[] i:f32[]. let         \n      ] a                                                                                     j:f32[] = sin h                           \n      f:f32[] = sin b                                                                         k:f32[] = cos h                           \n      g:f32[] = sin f                                                                         l:f32[] = cos j                           \n      h:f32[] = sin g                                                                         m:f32[] = mul i l                         \n      i:f32[] = sin h                                                                         n:f32[] = mul m k                         \n      j:f32[] = sin i                                                                       in (n,) }                                   \n      k:f32[] = cos i                                                                     policy=None                                   \n      l:f32[] = sin j                                                                     prevent_cse=True                              \n      m:f32[] = cos j                                                                   ] a f                                           \n    in (l, a, g, k, m) }                                                                o:f32[] = remat2[                               \n                                                                                          differentiated=True                           \n                                                                                          jaxpr={ lambda ; p:f32[] q:f32[]. let         \n                                                                                              r:f32[] = sin p                           \n                                                                                              s:f32[] = sin r                           \n                                                                                              t:f32[] = sin s                           \n                                                                                              u:f32[] = cos s                           \n                                                                                              v:f32[] = cos t                           \n                                                                                              w:f32[] = mul q v                         \n                                                                                              x:f32[] = mul w u                         \n                                                                                              y:f32[] = remat2[                         \n                                                                                                differentiated=True                     \n                                                                                                jaxpr={ lambda ; z:f32[] ba:f32[]. let  \n                                                                                                    bb:f32[] = sin z                    \n                                                                                                    bc:f32[] = cos z                    \n                                                                                                    bd:f32[] = cos bb                   \n                                                                                                    be:f32[] = mul ba bd                \n                                                                                                    bf:f32[] = mul be bc                \n                                                                                                  in (bf,) }                            \n                                                                                                policy=None                             \n                                                                                                prevent_cse=True                        \n                                                                                              ] p x                                     \n                                                                                            in (y,) }                                   \n                                                                                          policy=None                                   \n                                                                                          prevent_cse=True                              \n                                                                                        ] 3.0:f32[] g                                   \n                                                                                      in (o,) }\n```\n\nExample:\n```text\nLayerParam = tuple[jnp.ndarray, jnp.ndarray]  # Weights-bias pair for a layer.\nParamsList = list[LayerParam]\n\ndef net(params: ParamsList, x: jnp.ndarray):\n  for W, b in params:\n    x = jnp.maximum(jnp.dot(x, W) + b, 0.)\n  return x\n```\n\nExample:\n```text\nparams = [(jnp.array([[0.5, 0.5], [1., 1.]]), jnp.array([0.5, 0.5])),\n          (jnp.array([[0.5, 0.5], [1., 1.]]), jnp.array([0.5, 0.5]))]\n\nall_weights = jnp.stack([W for W, _ in params])\nall_biases = jnp.stack([b for _, b in params])\n\ndef layer(x, W_b_pair):\n  W, b = W_b_pair\n  out = jnp.maximum(jnp.dot(x, W) + b, 0.)\n  return out, None\n\ndef net(all_weights, all_biases, x):\n  x, _ = jax.lax.scan(layer, x, (all_weights, all_biases))\n  return x\n```\n\nExample:\n```text\nfrom functools import partial\n\n@partial(jax.checkpoint,\n         policy=jax.checkpoint_policies.dots_with_no_batch_dims_saveable)\ndef layer(x, W_b_pair):\n  W, b = W_b_pair\n  out = jnp.maximum(jnp.dot(x, W) + b, 0.)\n  return out, None\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.842Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":49,"totalLines":693,"estimatedTokens":8787}}119{"id":"doc-jep_28661_supporting_the_jax_array_protocol_jax_-e6f643f2","source":"documentation","title":"JEP 28661: Supporting the __jax_array__ protocol — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/28661-jax-array-protocol.html","text":"Example:\n```text\nclass CustomArray:\n  data: numpy.ndarray\n  ...\n\nx = CustomArray(np.arange(5))\nresult = jnp.sin(x)  # Converts `x` to JAX array and returns a JAX array\n```\n\nExample:\n```text\nclass CustomArray:\n  data: numpy.ndarray\n  ...\n\nx = CustomArray(np.arange(5))\nresult = jnp.sin(x)  # returns a new CustomArray\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\n\nclass CustomArray:\n  data: np.ndarray\n\n  def __init__(self, data: np.ndarray):\n    self.data = data\n\n  def __jax_array__(self) -> jax.Array:\n    return jnp.asarray(self.data)\n\narr = CustomArray(np.arange(5))\nresult = jnp.multiply(arr, 2)\nprint(repr(result))\n# Array([0, 2, 4, 6, 8], dtype=int32)\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  print(\"is JAX array:\", isinstance(x, jax.Array))\n\nf(CustomArray(...))\n```\n\nExample:\n```text\nis JAX array: True\n```\n\nExample:\n```text\ntype(x)=CustomArray\n```\n\nExample:\n```text\ndef f(x):\n  if isinstance(x, CustomArray):\n    return x.custom_method()\n  else:\n    # do something else\n    ...\n\nresult1 = f(x)\nresult2 = jax.jit(f)(x)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.843Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":7,"totalLines":74,"estimatedTokens":270}}120{"id":"doc-jax_and_jaxlib_versioning_jax_documentation-67096baa","source":"documentation","title":"Jax and Jaxlib versioning — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/9419-jax-versioning.html","text":"Example:\n```text\nfrom jax._src.lib import jaxlib_extension_version\n\n# 123 is the new version number for _version in xla_client.py\nif jaxlib_extension_version >= 123:\n  # Use new code path\n  ...\nelse:\n  # Use old code path.\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.845Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":1,"totalLines":13,"estimatedTokens":60}}121{"id":"doc-jax_extend_a_module_for_extensions_jax_documenta-3ea45bb8","source":"documentation","title":"jax.extend: a module for extensions — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/15856-jex.html","text":"Example:\n```text\nimport jax.extend as jex\n```\n\nExample:\n```text\nfrom jax.extend import core\t         # Previously: from jax import core\nfrom jax.extend.interpreters import mlir        # ... and similarly\n\nmul_add_p = core.Primitive('mul_add')\nmul_add_p.def_impl(lambda x, y, z: x * y + z)\n\n@mul_add_p.def_abstract_eval\ndef mul_add_abstract(x_sa, y_sa, z_sa):\n  return core.ShapedArray(x_sa.shape, x_sa.dtype)\n\ndef mul_add_mlir(ctx, xc, yc, zc):\n  add = mlir.hlo.AddOp\n  mul = mlir.hlo.MulOp\n  return add(mul(xc, yc), zc).results\n\nmlir.register_lowering(mul_add_p, mul_add_mlir)\n\nimport jax\nprint(mul_add_p.bind(2, 3, 4))            # -> 10\nprint(jax.jit(mul_add_p.bind)(2, 3, 4))   # -> Array(10, dtype=int32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.846Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":2,"totalLines":30,"estimatedTokens":182}}122{"id":"doc-sequencing_side_effects_in_jax_jax_documentation-2358dd55","source":"documentation","title":"Sequencing side-effects in JAX — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/10657-sequencing-effects.html","text":"Example:\n```text\ndef f():\n  print(\"hello\")\n  return 2\ndef g():\n  print(\"world\")\n  return 3\nf()\ng()\n```\n\nExample:\n```text\n@jax.jit(device=<device 0>)\ndef f():\n  return 2\n\n@jax.jit(device=<device 1>)\ndef g():\n  return 3\nf()\ng()\n```\n\nExample:\n```text\n@jax.jit(device=<device 0>)\ndef f():\n  jax.print(\"hello\")\n  return 2\n\n@jax.jit(device=<device 1>)\ndef g():\n  jax.print(\"world\")\n  return 3\nf()\ng()\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  jax.print(\"hello\")\n  jax.print(\"world\")\n  return x\n```\n\nExample:\n```text\n@jax.jit\ndef f(x, y):\n  log_value(x)\n  log_value(y)\nf(1, 2)\n```\n\nExample:\n```text\n@jax.jit\ndef f(token, x):\n  token = jax.print(token, \"hello\")\n  token = jax.print(token, \"world\")\n  return token, x\n```\n\nExample:\n```text\n@jax.jit\ndef f(runtime_token, x):\n  compiler_token = new_compiler_token()\n  compiler_token = jax.print(compiler_token, \"hello\")\n  compiler_token = jax.print(compiler_token, \"world\")\n  return runtime_token, x\n```\n\nExample:\n```text\ndef _execute(compiled_computation, *args):\n  outputs = compiled_computation.execute(*args)\n  return outputs\n```\n\nExample:\n```text\ndef _execute(compiled_computation, *args):\n  runtime_token = get_runtime_token() # Grab global token\n  runtime_token, *outputs = compiled_computation.execute(runtime_token, *args)\n  update_runtime_token(runtime_token) # Update global token\n  return outputs\n```\n\nExample:\n```text\n@jax.jit\ndef f():\n  jax.print(\"hello world\")\n  return\nf() # Executed asynchronously\n```\n\nExample:\n```text\n@jax.jit\ndef f():\n  jax.print(\"hello world\")\n  return new_runtime_token()\nf() # Executed asynchronously\n```\n\nExample:\n```text\n@jax.jit\ndef f():\n  jax.print(\"hello\")\n\n@jax.jit\ndef g():\n  jax.print(\"world\")\n\nf()\ng()\n```\n\nExample:\n```text\n@jax.jit(device=<device 0>)\ndef f():\n  jax.print(\"hello\")\n\n@jax.jit(device=<device 1>)\ndef g():\n  jax.print(\"world\")\n\nf()\ng()\n```\n\nExample:\n```text\n@jax.jit(device=<device 0>)\ndef f():\n  jax.print(\"hello\")\n  return new_runtime_token()\n\n@jax.jit(device=<device 1>)\ndef g():\n  jax.print(\"world\")\n  return new_runtime_token()\n\nt0 = f()\nt1 = g()\nblock_until_ready((t0, t1))\n```\n\nExample:\n```text\ndef _execute(compiled_computation, *args):\n  output_token, *outputs = compiled_computation.execute(runtime_token, *args)\n  update_output_token(output_token, compiled_computation.device)\n  return outputs\n```\n\nExample:\n```text\ndef effects_barrier():\n  output_token.block_until_ready()\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.848Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":16,"totalLines":171,"estimatedTokens":601}}123{"id":"doc-jep_28845_stateful_randomness_in_jax_jax_documen-a2eabb03","source":"documentation","title":"JEP 28845: Stateful Randomness in JAX — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/28845-stateful-rng.html","text":"Example:\n```text\ndef stateful_rng(seed: ArrayLike) -> StatefulPRNG:\n  \"\"\"Create a stateful PRNG Generator given an integer seed.\"\"\"\n  return StatefulPRNG(jax.random.key(seed), jax.new_ref(0))\n\n\n@tree_util.register_dataclass\n@dataclass(frozen=True)\nclass StatefulPRNG:\n  \"\"\"Stateful PRNG Generator class.\"\"\"\n  base_key: jax.Array\n  counter: jax.core.Ref\n\n  def key(self) -> jax.Array:\n    \"\"\"Generate a new jax PRNG key\"\"\"\n    key = jax.random.fold_in(self.base_key, self.counter[...])\n    jax.ref.addupdate(self.counter, ..., 1)  # increment counter\n    return key\n\n  def random(self, size: Sequence[int], dtype: DType = float):\n    \"\"\"Return random floats in the half-open interval [0, 1)\"\"\"\n    return random.uniform(self.key(), shape=size, dtype=dtype)\n\n  # uniform(), normal(), integers(), and others implemented similarly.\n```\n\nExample:\n```text\n>>> from jax.experimental.random import stateful_rng\n>>> rng = stateful_rng(1701)\n>>> rng.random((5,))\nArray([0.09609699, 0.26730824, 0.5619041 , 0.24421775, 0.7715055 ], dtype=float32)\n>>> rng.random((5,))  # state is updated -> new random draws!\nArray([0.8131045 , 0.33873856, 0.88808906, 0.96005905, 0.7616446 ], dtype=float32)\n\n>>> import numpy as np\n>>> rng = np.random.default_rng(1701)\n>>> rng.random((5,))\narray([0.4020733 , 0.30563311, 0.67668051, 0.15821208, 0.79247763])\n>>> rng.random((5,))\narray([0.09419469, 0.36753944, 0.06388928, 0.96431608, 0.35200998])\n```\n\nExample:\n```text\nrng = stateful_rng(0)\n\ndef f(x):\n  return x + rng.uniform()\n\njax.vmap(f)(jnp.arange(10))\n```\n\nExample:\n```text\nException: performing an addupdate operation with vmapped value on an unbatched\n           array reference of type Ref{int32[]}. Move the array reference to be\n           an argument to the vmapped function?\n```\n\nExample:\n```text\nclass StatefulPRNG:\n  ...\n\n  def split(self, num: int | Sequence[int]) -> StatefulPRNG:\n    return StatefulPRNG(\n      base_key=jax.random.split(self.key(), num),\n      counter=jnp.zeros(num, dtype=int),\n    )\n```\n\nExample:\n```text\nrng = jax.experimental.random.stateful_rng(0)\n\ndef f(x, rng):\n  return x + rng.uniform()\n\nresult = jax.vmap(f)(jnp.arange(5), rng.split(5))\nprint(result)  # [0.07174575 1.0163325  2.0435536  3.4391735  4.534091  ]\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.849Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":6,"totalLines":85,"estimatedTokens":562}}124{"id":"doc-jep_9263_typed_keys_pluggable_rngs_jax_documenta-f36255bc","source":"documentation","title":"JEP 9263: Typed keys & pluggable RNGs — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/9263-typed-keys.html","text":"Example:\n```text\n>>> key = jax.random.PRNGKey(0)\n>>> key\nArray([0, 0], dtype=uint32)\n>>> key.shape\n(2,)\n>>> key.dtype\ndtype('uint32')\n```\n\nExample:\n```text\n>>> key = jax.random.key(0)\n>>> key\nArray((), dtype=key<fry>) overlaying:\n[0 0]\n>>> key.shape\n()\n>>> key.dtype\nkey<fry>\n```\n\nExample:\n```text\n>>> key_arr = jax.vmap(jax.random.key)(jnp.arange(4))\n>>> key_arr\nArray((4,), dtype=key<fry>) overlaying:\n[[0 0]\n [0 1]\n [0 2]\n [0 3]]\n>>> key_arr.shape\n(4,)\n```\n\nExample:\n```text\n# split\nnew_key, subkey = jax.random.split(key)\n\n# random number generation\ndata = jax.random.uniform(key, shape=(5,))\n```\n\nExample:\n```text\n>>> key = key + 1  \nTraceback (most recent call last):\nTypeError: add does not accept dtypes key<fry>, int32.\n```\n\nExample:\n```text\n>>> jax.random.key_data(key)\nArray([0, 0], dtype=uint32)\n```\n\nExample:\n```text\n>>> typed_key = jax.random.key(0)\n>>> jax.dtypes.issubdtype(typed_key.dtype, jax.dtypes.prng_key)\nTrue\n>>> raw_key = jax.random.PRNGKey(0)\n>>> jax.dtypes.issubdtype(raw_key.dtype, jax.dtypes.prng_key)\nFalse\n```\n\nExample:\n```text\nfrom jax import dtypes\n\ndef ensure_typed_key_array(key: Array) -> Array:\n  if dtypes.issubdtype(key.dtype, dtypes.prng_key):\n    return key\n  else:\n    raise TypeError(\"New-style typed JAX PRNG keys required\")\n```\n\nExample:\n```text\n>>> key = jax.random.key(0, impl='threefry2x32')  # this is the default impl\n>>> key\nArray((), dtype=key<fry>) overlaying:\n[0 0]\n>>> jax.random.uniform(key, shape=(3,))\nArray([0.947667  , 0.9785799 , 0.33229148], dtype=float32)\n\n>>> key = jax.random.key(0, impl='rbg')\n>>> key\nArray((), dtype=key<rbg>) overlaying:\n[0 0 0 0]\n>>> jax.random.uniform(key, shape=(3,))\nArray([0.39904642, 0.8805201 , 0.73571277], dtype=float32)\n```\n\nExample:\n```text\n# Incorrect\nkey = random.PRNGKey(999)\nnew_key = random.PRNGKey(key[1])  # identical to the original key!\n```\n\nExample:\n```text\n# Correct\nkey = random.PRNGKey(999)\nkey, new_key = random.split(key)\n```\n\nExample:\n```text\n# Incorrect\nkey = random.PRNGKey(0)\nbatched_keys = key + jnp.arange(10, dtype=key.dtype)[:, None]\n```\n\nExample:\n```text\n# Correct\nkey = random.PRNGKey(0)\nbatched_keys = random.split(key, 10)\n```\n\nExample:\n```text\n# Incorrect\nkeys = random.split(random.PRNGKey(0))\ndata = jax.vmap(random.uniform, in_axes=1)(keys)\n```\n\nExample:\n```text\n# Correct\nkeys = random.split(random.PRNGKey(0))\ndata = jax.vmap(random.uniform, in_axes=0)(keys)\n```\n\nExample:\n```text\n# Incorrect\nkey = random.PRNGKey(0)\nx = random.uniform(key, (100,))\ny = random.uniform(key, (100,))  # Identical values!\n```\n\nExample:\n```text\n# Correct\nkey = random.PRNGKey(0)\nkey1, key2 = random.split(random.key(0))\nx = random.uniform(key1, (100,))\ny = random.uniform(key2, (100,))\n```\n\nExample:\n```text\n>>> jax.dtypes.issubdtype(jax.dtypes.prng_key, jax.dtypes.extended)\nTrue\n```\n\nExample:\n```text\n>>> key = jax.random.key(0)\n>>> jax.dtypes.issubdtype(key.dtype, jax.dtypes.extended)\nTrue\n>>> jax.dtypes.issubdtype(key.dtype, jax.dtypes.prng_key)\nTrue\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.850Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":19,"totalLines":171,"estimatedTokens":746}}125{"id":"doc-type_annotation_roadmap_for_jax_jax_documentatio-61fe41f4","source":"documentation","title":"Type Annotation Roadmap for JAX — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/12049-type-annotations.html","text":"Example:\n```text\nArray = Any\nShape = core.Shape\n\ndef slice(operand: Array, start_indices: Sequence[int],\n          limit_indices: Sequence[int],\n          strides: Optional[Sequence[int]] = None) -> Array:\n  ...\n```\n\nExample:\n```text\ndef shuffle(key: KeyArray, x: Array, axis: int = 0) -> jnp.ndarray:\n  ...\n```\n\nExample:\n```text\ndef tile(A, reps):\n  try:\n    tup = tuple(reps)\n  except TypeError:\n    tup = (reps,)\n  d = len(tup)\n  ...\n```\n\nExample:\n```text\n@jit\ndef f(x: ArrayAnnotation) -> ArrayAnnotation:\n  assert isinstance(x, core.Tracer)\n  return x\n```\n\nExample:\n```text\ndef f(x):\n  return isinstance(x, ArrayInstance)\nx = jnp.array([1, 2, 3])\nassert f(x)       # x will be an array\nassert jit(f)(x)  # x will be a tracer\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.853Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":5,"totalLines":46,"estimatedTokens":187}}126{"id":"doc-shmap_shard_map_for_simple_per_device_code_jax_d-f2d67681","source":"documentation","title":"shmap (shard_map) for simple per-device code — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/14273-shard-map.html","text":"Example:\n```text\nfrom functools import partial\n\nimport numpy as np\n\nimport jax\nimport jax.numpy as jnp\nfrom jax.sharding import Mesh, PartitionSpec as P\nfrom jax.experimental.shard_map import shard_map\n\nmesh = jax.make_mesh((4, 2), ('i', 'j'))\n\na = jnp.arange( 8 * 16.).reshape(8, 16)\nb = jnp.arange(16 * 32.).reshape(16, 32)\n\n@partial(shard_map, mesh=mesh, in_specs=(P('i', 'j'), P('j', None)),\n         out_specs=P('i', None))\ndef matmul_basic(a_block, b_block):\n  # a_block: f32[2, 8]\n  # b_block: f32[8, 32]\n  z_partialsum = jnp.dot(a_block, b_block)\n  z_block = jax.lax.psum(z_partialsum, 'j')\n  return z_block\n\nc = matmul_basic(a, b)  # c: f32[8, 32]\n```\n\nExample:\n```text\n@partial(shard_map, mesh=mesh, in_specs=(P('i', 'j'), P('j', None)),\n         out_specs=P('i', 'j'))\ndef matmul_reduce_scatter(a_block, b_block):\n  # c_partialsum: f32[8/X, 32]\n  c_partialsum = jnp.matmul(a_block, b_block)\n  # c_block: f32[8/X, 32/Y]\n  c_block = jax.lax.psum_scatter(c_partialsum, 'j', scatter_dimension=1, tiled=True)\n  return c_block\n\nc = matmul_reduce_scatter(a, b)\n```\n\nExample:\n```text\npmap(f, in_axes=[0], out_axes=0)(xs) == jnp.stack([f(x) for x in xs])\n```\n\nExample:\n```text\ndevices = np.array(jax.devices()[:4])\nm = Mesh(devices, ('i',))  # mesh.shape['i'] = 4\n\nshard_map(f, m, in_specs=P('i'), out_specs=P('i'))(y)\n==\njnp.concatenate([f(y_blk) for y_blk in jnp.split(y, 4)])\n```\n\nExample:\n```text\ndevices = np.array(jax.devices())\nm = Mesh(devices.reshape(4, 2), ('i', 'j'))\n\n@partial(shard_map, mesh=m, in_specs=P('i', None), out_specs=P('i', 'j'))\ndef f1(x_block):\n  print(x_block.shape)\n  return x_block\n\nx1 = np.arange(12 * 12).reshape(12, 12)\ny = f1(x1)  # prints (3,12)\n```\n\nExample:\n```text\n@partial(shard_map, mesh=m, in_specs=P('i', 'j'), out_specs=P('i', 'j'))\ndef f2(x_block):\n  print(x_block.shape)\n  return x_block\n\nx = np.arange(12 * 12).reshape(12, 12)\nx_ = jnp.tile(x, (1, mesh.axis_size['j']))  # x_ has shape (12, 24)\ny = f2(x_)  # prints (3,12), and f1(x) == f2(x_)\n```\n\nExample:\n```text\nx = jnp.array([[3.]])\n\nz = shard_map(lambda: x, mesh=m, in_specs=(), out_specs=P('i', 'j'))()\nprint(z)  # prints the same as jnp.tile(x, (4, 2))\n\nz = shard_map(lambda: x, mesh=m, in_specs=(), out_specs=P('i', None))()\nprint(z)  # prints the same as jnp.tile(x, (4, 1)), or just jnp.tile(x, (4,))\n\nz = shard_map(lambda: x, mesh=m, in_specs=(), out_specs=P(None, None))()\nprint(z)  # prints the same as jnp.tile(x, (1, 1)), or just x\n```\n\nExample:\n```text\n@partial(shard_map, mesh=m, in_specs=P('i', 'j'), out_specs=P('i', None))\ndef f3(x_block):\n  return jax.lax.psum(x_block, 'j')\n\nx = np.arange(12 * 12).reshape(12, 12)\ny3 = f3(x)\nprint(y3.shape)  # (12,6)\n```\n\nExample:\n```text\n@partial(shard_map, mesh=m, in_specs=P('i', 'j'), out_specs=P(None, 'j'))\ndef f4(x_block):\n  return jax.lax.psum(x_block, 'i')\n\nx = np.arange(12 * 12).reshape(12, 12)\ny4 = f4(x)\nprint(y4.shape)  # (3,12)\n\n\n@partial(shard_map, mesh=m, in_specs=P('i', 'j'), out_specs=P(None, None))\ndef f5(x_block):\n  return jax.lax.psum(x_block, ('i', 'j'))\n\ny5 = f5(x)\nprint(y5.shape)  # (3,6)\n```\n\nExample:\n```text\nfrom jax.sharding import Mesh\nSpecs = PyTree[PartitionSpec]\n\ndef shard_map(f: Callable, mesh: Mesh, in_specs: Specs, out_specs: Specs\n          ) -> Callable:\n  ...\n```\n\nExample:\n```text\ndef matmul_2D_wg_manual(xnorm, q_wi, layer):\n  '''Calls a custom manual implementation of matmul_reducescatter'''\n  # [batch, maxlen, embed.X] @ [heads.YZ, embed.X, q_wi_per_head]\n  # -> (matmul)\n  # -> [batch, maxlen, heads.YZ, q_wi_per_head]{x unreduced}\n  # -> (reducescatter over x into X heads, B batches)\n  # -> [batch, maxlen, heads.YZX, q_wi_per_head]\n  with jax.named_scope('q_wi'):\n    xnorm = intermediate_dtype(xnorm)\n    q_wi = matmul_reducescatter(\n        'bte,hed->bthd',\n        xnorm,\n        params.q_wi,\n        scatter_dimension=(0, 2),\n        axis_name='i',\n        layer=layer)\n   return q_wi\n\n\nimport partitioning.logical_to_physical as l2phys\n\ndef pjit_transformer_layer(\n    hparams: HParams, layer: int, params: weights.Layer, sin: jnp.ndarray,\n    cos: jnp.ndarray, kv_caches: Sequence[attention.KVCache],\n    x: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:\n  \"\"\"Forward pass through a single layer, returning output, K, V.\"\"\"\n\n  def my_layer(t, axis=0):\n    \"\"\"Gets the parameters corresponding to a given layer.\"\"\"\n    return lax.dynamic_index_in_dim(t, layer, axis=axis, keepdims=False)\n\n  # 2D: [batch.Z, time, embed.XY]\n  x = _with_sharding_constraint(\n      x, ('residual_batch', 'residual_time', 'residual_embed'))\n  xnorm = _layernorm(x)\n  # 2D: [batch, time, embed.X]\n  xnorm = _with_sharding_constraint(\n      xnorm, ('post_norm_batch', 'time', 'post_norm_embed'))\n  # jump into manual mode where you want to optimise\n  if manual:\n    q_wi = shard_map(matmul_2D_wg_manual, mesh\n                in_specs=(l2phys('post_norm_batch', 'time', 'post_norm_embed'),\n                          l2phys('layers', 'heads', 'embed', 'q_wi_per_head')),\n                out_specs=l2phys('post_norm_batch', 'time', 'heads', 'q_wi_per_head'))(xnorm, q_wi, layer)\n  else:\n    q_wi = jnp.einsum('bte,hed->bthd', xnorm, my_layer(params.q_wi))\n    # 2D: [batch, time, heads.YZX, None]\n    q_wi = _with_sharding_constraint(q_wi,\n                                   ('post_norm_batch', 'time', 'heads', 'qkv'))\n  q = q_wi[:, :, :, :hparams.qkv]\n  q = _rope(sin, cos, q)\n  # unlike in https://arxiv.org/pdf/2002.05202.pdf, PaLM implements\n  # swiGLU with full d_ff dimension, rather than 2/3 scaled\n  wi0 = q_wi[:, :, :, hparams.qkv:hparams.qkv + (hparams.ff // hparams.heads)]\n  wi1 = q_wi[:, :, :, hparams.qkv + (hparams.ff // hparams.heads):]\n  kv = jnp.einsum('bte,ezd->btzd', xnorm, my_layer(params.kv))\n  k = kv[:, :, 0, :hparams.qkv]\n  v = kv[:, :, 0, hparams.qkv:]\n  k = _rope(sin, cos, k)\n\n  y_att = jnp.bfloat16(attention.attend(q, k, v, kv_caches, layer))\n\n  y_mlp = special2.swish2(wi0) * wi1\n  # 2D: [batch, time, heads.YZX, None]\n  y_mlp = _with_sharding_constraint(y_mlp,\n                                    ('post_norm_batch', 'time', 'heads', None))\n\n  y_fused = jnp.concatenate([y_att, y_mlp], axis=-1)\n  # do the second half of the mlp and the self-attn projection in parallel\n  y_out = jnp.einsum('bthd,hde->bte', y_fused, my_layer(params.o_wo))\n  # 2D: [batch.Z, time, embed.XY]\n  y_out = _with_sharding_constraint(\n      y_out, ('residual_batch', 'residual_time', 'residual_embed'))\n  z = y_out + x\n  z = _with_sharding_constraint(\n      z, ('residual_batch', 'residual_time', 'residual_embed'))\n  return z, k, v\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.854Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":11,"totalLines":219,"estimatedTokens":1641}}127{"id":"doc-efficient_transposition_of_replication_inducing_-5fbf8b8a","source":"documentation","title":"Efficient transposition of replication-inducing collectives — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/17111-shmap-transpose.html","text":"Example:\n```text\ndevices = jax.devices()  # 8 devices\n\n@partial(shmap, mesh=Mesh(devices, ('batch',)),\n         in_specs=(P(None, None), P('batch', None)),\n         out_specs=P())\ndef loss(params, batch):\n  inputs, targets = batch\n  predictions = predict(params, inputs)\n  local_loss = jnp.mean(jnp.sum(predictions - targets, -1))\n  global_loss = lax.pmean(local_loss, 'batch'))\n  return global_loss\n```\n\nExample:\n```text\n# Example 1: shmap involving psum and unmapped output with inefficient transpose\nf1 = shmap(lambda x: psum(g(x), 'i'),\n           in_specs=P('i'), out_specs=P())\n```\n\nExample:\n```text\n# An efficient \"transpose\" of Example 1 (but don't transpose this again!)\n¿f1_transpose? = shmap(t(g), in_specs=P(), out_specs=P('i'))\n```\n\nExample:\n```text\n# The transpose we currently get for Example 1 (which is fine to transpose again)\nt(f1) = shmap(lambda ybar: t(g)(psum(ybar / 8, 'i')),\n              in_specs=P(), out_specs=P('i'))\n```\n\nExample:\n```text\n# Example 2: shmap involving psum and *mapped* output with efficient transpose\nf2 = shmap(lambda x, y: psum(g(x), 'i') * y,\n          in_specs=(P('i'), P('i')), out_specs=P('i'))\n\n# The transpose we currently get for Example 2 is efficient\nt(f2, 0) = shmap(lambda y, zbar: t(g)(psum(zbar * y, 'i')),\n                in_specs=(P('i'), P('i')), out_specs=P('i'))\n```\n\nExample:\n```text\n# Example 3: cursed identity\ncursed_identity = shmap(lambda x: x, P(), P())\n\n# Currently we get these inefficient transposes\nt(cursed_identity) = shmap(lambda x: psum(x / 8, 'i'), P(), P())\nt(t(cursed_identity)) = shmap(lambda x: psum(psum(x / 8 / 8, 'i'), 'i')), P(), P())\n...\n```\n\nExample:\n```text\n# Example 4: all_gather to an unmapped output\nf4 = shmap(lambda x: all_gather(x, 'i'), P('i'), P())\n\n# Currently we get this inefficient transpose\nt(f4) = shmap(lambda ybar: psum_scatter(ybar / 8, 'i'), P(), P('i'))\n```\n\nExample:\n```text\n# Example 5: all_gather to a mapped output\nf5 = shmap(lambda x, y: all_gather(x, 'i') * y,\n           in_specs=(P('i'), P('i')), out_specs=P('i'))\n\n# Currently we get this efficient transpose\nt(f5, 0) = shmap(lambda y, zbar: psum_scatter(zbar * y, 'i'),\n                 in_specs=(P('i'), P('i')), out_specs=P('i'))\n```\n\nExample:\n```text\n# Example 4 again\nf4 = shmap(lambda x: all_gather(x, 'i'), P('i'), P())\n\n# Why didn't we just write it like this?\nf4_better = shmap(lambda x: x, P('i'), P('i'))\n```\n\nExample:\n```text\n# Example 1 again\nf1 = shmap(lambda x: psum(g(x), 'i'),\n           in_specs=P('i'), out_specs=P())\n\n# What if we could write an output sum like this?\nf1_better = shmap(g, in_specs=P('i'), out_specs=P(sum='i'))  # sum='i' means sum over that axis\n\n# Then it could transpose like this:\nt(f1_better) = shmap(t(g), in_specs=P(), out_specs=P('i'))\nt(t(f1_better)) = shmap(t(t(g)), in_specs=P('i'), P(sum='i'))\n```\n\nExample:\n```text\n# Example 3 again\ncursed_identity = shmap(lambda x: x, P(), P())\n\n# How it would transpose with the P-sum partial solution:\nt(cursed_identity) = shmap(lambda x: x / 8, P(), P(sum='i'))\nt(t(cursed_identity)) = shmap(lambda x: x / 8, P(), P(sum='i'))\n```\n\nExample:\n```text\nshaped_array ::= <dtype>[<int_literal>, ...]<device_variance_type>\ndevice_variance_type ::= {<axis_name>, ...}\n```\n\nExample:\n```text\n# Example 1 again\nf1 = shmap(lambda x: psum(g(x), 'i'),\n           in_specs=P('i'), out_specs=P())\n\n# Example 1 with intermediate device variance types annotated\n@partial(shmap, in_specs=P('i'), out_specs=P())\ndef f1(x: f32[3,4]{i}):\n  w:f32[]{i} = g(x)\n  y:f32[]{} = psum(w, 'i')\n  return y\n```\n\nExample:\n```text\n# Example 1 transpose using device variance types (go ahead and transpose this again!)\nt(f1) = shmap(lambda ybar: t(g)(pbroadcast(ybar, 'i')),\n              in_specs=P(), out_specs=P('i'))\n\n# Example 1 transpose with intermediate device variance types annotated\n@partial(shmap, in_specs=P('i'), out_specs=P())\ndef f1_transpose(ybar: f32[]):\n  wbar:f32[]{i} = pbroadcast(ybar, 'i')\n  xbar:f32[3,4]{i} = transpose(g)(wbar)\n  return xbar\n```\n\nExample:\n```text\n# Example 2 rewritten with explicit pbroadcast\nf2 = shmap(lambda x, y: pbroadcast(psum(g(x), 'i'), 'i') * y,\n           in_specs=(P('i'), P('i')), out_specs=P('i'))\n\n# Example 2 transpose using device variance types\nt(f2, 0) = shmap(lambda y, zbar: t(g)(pbroadcast(psum(zbar * y, 'i'), 'i')),\n                 in_specs=(P('i'), P('i')), out_specs=P('i'))\n\n\n# Example 3 again\ncursed_identity = shmap(lambda x: x, P(), P())\n# Notice here the body is `f32[...] -> f32[...]`, i.e. no device varying type.\n\n# Example 3 transpose using device variance types\nt(cursed_identity) = shmap(lambda x: x, P(), P())\nt(t(cursed_identity)) = shmap(lambda x: x, P(), P())\n```\n\nExample:\n```text\n# Example 4 rewritten with explicit all_reduce_invariant\nf4 = shmap(lambda x: all_gather_invariant(x, 'i'), P('i'), P())\n\n# Example 4 with intermediate device variance types annotated\n@partial(shmap, P('i'), P())\ndef f4(x:f32[1]{i}):\n  y:f32[8]{} = all_gather_invariant(x, 'i')\n  return y\n\n# Example 4 transpose with intermediate device variance types annotated\n@partial(shmap, in_specs=P(), out_specs=P('i'))\ndef f4_transpose(ybar:f32[8]):\n  xbar:f32[1]{i} = pscatter(ybar, 'i')\n  return xbar\n```\n\nExample:\n```text\n# Example 5 with intermediate device variance types annotated\n@partial(shmap, in_specs=(P('i'), P('i')), out_specs=P('i'))\ndef f5(x:f32[1]{i}, y:f32[8]{i}):\n  z:f32[8]{i} = all_gather(x, 'i')\n  w:f32[8]{i} = z * y\n  return w\n\n# Transpose with respect to first argument\n@partial(shmap, in_specs=(P('i'), P('i')), out_specs=P('i'))\ndef f5_transpose(y:f32[8]{i}, wbar:f32[8]{i}):\n  zbar:f32[8]{i} = wbar * y\n  xbar:f32[1]{i} = psum_scatter(zbar, 'i')\n  return xbar\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.856Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":17,"totalLines":200,"estimatedTokens":1421}}128{"id":"doc-default_dtypes_and_the_x64_flag_jax_documentatio-9811f964","source":"documentation","title":"Default dtypes and the X64 flag — JAX documentation","url":"https://docs.jax.dev/en/latest/default_dtypes.html","text":"Example:\n```text\n>>> import jax.numpy as jnp\n\n>>> jnp.arange(5)\nArray([0, 1, 2, 3, 4], dtype=int32)\n\n>>> jnp.zeros(5)\nArray([0., 0., 0., 0., 0.], dtype=float32)\n\n>>> jnp.ones(5, dtype=int)\nArray([1, 1, 1, 1, 1], dtype=int32)\n```\n\nExample:\n```text\n>>> jnp.arange(5, dtype='float64')  \nUserWarning: Explicitly requested dtype float64 requested in arange is not available, and will be \ntruncated to dtype float32. To enable more dtypes, set the jax_enable_x64 configuration option or the \nJAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.\nArray([0., 1., 2., 3., 4.], dtype=float32)\n```\n\nExample:\n```text\nimport jax\nimport jax.numpy as jnp\n\njax.config.update('jax_enable_x64', True)\n\nprint(repr(jnp.arange(5)))\nprint(repr(jnp.zeros(5)))\nprint(repr(jnp.ones(5, dtype=int)))\n```\n\nExample:\n```text\nArray([0, 1, 2, 3, 4], dtype=int64)\nArray([0., 0., 0., 0., 0.], dtype=float64)\nArray([1, 1, 1, 1, 1], dtype=int64)\n```\n\nExample:\n```text\n$ JAX_ENABLE_X64=1 python main.py\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.856Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":5,"totalLines":48,"estimatedTokens":258}}129{"id":"doc-asynchronous_dispatch_jax_documentation-5f85784c","source":"documentation","title":"Asynchronous dispatch — JAX documentation","url":"https://docs.jax.dev/en/latest/async_dispatch.html","text":"Example:\n```text\n>>> import numpy as np\n>>> import jax.numpy as jnp\n>>> from jax import random\n>>> x = random.uniform(random.key(0), (1000, 1000))\n>>> # Printing the result (i.e. evaluating `repr(result)` or `str(result)`)\n>>> # will block until the value is ready.\n>>> jnp.dot(x, x) + 3.  \nArray([[258.01971436, 249.64862061, 257.13372803, ...,\n        236.67948914, 250.68939209, 241.36853027],\n        [265.65979004, 256.28912354, 262.18252563, ...,\n        242.03181458, 256.16757202, 252.44122314],\n        [262.38916016, 255.72747803, 261.23059082, ...,\n        240.83563232, 255.41094971, 249.62471008],\n        ...,\n        [259.15814209, 253.09197998, 257.72174072, ...,\n        242.23876953, 250.72680664, 247.16642761],\n        [271.22662354, 261.91204834, 265.33398438, ...,\n        248.26651001, 262.05389404, 261.33700562],\n        [257.16134644, 254.7543335, 259.08300781, ..., 241.59848022,\n        248.62597656, 243.22348022]], dtype=float32)\n```\n\nExample:\n```text\n>>> %time jnp.dot(x, x)  \nCPU times: user 267 µs, sys: 93 µs, total: 360 µs\nWall time: 269 µs\nArray([[255.01972961, 246.64862061, 254.13371277, ...,\n        233.67948914, 247.68939209, 238.36853027],\n        [262.65979004, 253.28910828, 259.18252563, ...,\n        239.03181458, 253.16757202, 249.44122314],\n        [259.38916016, 252.72747803, 258.23059082, ...,\n        237.83563232, 252.41094971, 246.62471008],\n        ...,\n        [256.15814209, 250.09197998, 254.72172546, ...,\n        239.23876953, 247.72680664, 244.16642761],\n        [268.22662354, 258.91204834, 262.33398438, ...,\n        245.26651001, 259.05389404, 258.33700562],\n        [254.16134644, 251.7543335, 256.08300781, ..., 238.59848022,\n        245.62597656, 240.22348022]], dtype=float32)\n```\n\nExample:\n```text\n>>> %time np.asarray(jnp.dot(x, x))  \nCPU times: user 61.1 ms, sys: 0 ns, total: 61.1 ms\nWall time: 8.09 ms\nOut[16]:\narray([[255.01973, 246.64862, 254.13371, ..., 233.67949, 247.68939,\n        238.36853],\n       [262.6598 , 253.28911, 259.18253, ..., 239.03181, 253.16757,\n        249.44122],\n       [259.38916, 252.72748, 258.2306 , ..., 237.83563, 252.41095,\n        246.62471],\n       ...,\n       [256.15814, 250.09198, 254.72173, ..., 239.23877, 247.7268 ,\n        244.16643],\n       [268.22662, 258.91205, 262.33398, ..., 245.26651, 259.0539 ,\n        258.337  ],\n       [254.16135, 251.75433, 256.083  , ..., 238.59848, 245.62598,\n        240.22348]], dtype=float32)\n>>> %time jnp.dot(x, x).block_until_ready()  \nCPU times: user 50.3 ms, sys: 928 µs, total: 51.2 ms\nWall time: 4.92 ms\nArray([[255.01972961, 246.64862061, 254.13371277, ...,\n        233.67948914, 247.68939209, 238.36853027],\n        [262.65979004, 253.28910828, 259.18252563, ...,\n        239.03181458, 253.16757202, 249.44122314],\n        [259.38916016, 252.72747803, 258.23059082, ...,\n        237.83563232, 252.41094971, 246.62471008],\n        ...,\n        [256.15814209, 250.09197998, 254.72172546, ...,\n        239.23876953, 247.72680664, 244.16642761],\n        [268.22662354, 258.91204834, 262.33398438, ...,\n        245.26651001, 259.05389404, 258.33700562],\n        [254.16134644, 251.7543335, 256.08300781, ..., 238.59848022,\n        245.62597656, 240.22348022]], dtype=float32)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.858Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":82,"estimatedTokens":811}}130{"id":"doc-rank_promotion_warning_jax_documentation-0aac9f19","source":"documentation","title":"Rank promotion warning — JAX documentation","url":"https://docs.jax.dev/en/latest/rank_promotion_warning.html","text":"Example:\n```text\n>>> from jax import numpy as jnp\n>>> x = jnp.arange(12).reshape(4, 3)\n>>> y = jnp.array([0, 1, 0])\n>>> x + y\nArray([[ 0,  2,  2],\n       [ 3,  5,  5],\n       [ 6,  8,  8],\n       [ 9, 11, 11]], dtype=int32)\n```\n\nExample:\n```text\nwith jax.numpy_rank_promotion(\"warn\"):\n  z = x + y\n```\n\nExample:\n```text\nimport jax\njax.config.update(\"jax_numpy_rank_promotion\", \"warn\")\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.859Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":3,"totalLines":25,"estimatedTokens":101}}131{"id":"doc-frequently_asked_questions_faq_jax_documentation-ea5f21d6","source":"documentation","title":"Frequently asked questions (FAQ) — JAX documentation","url":"https://docs.jax.dev/en/latest/faq.html","text":"Example:\n```text\ny = 0\n\n# @jit   # Different behavior with jit\ndef impure_func(x):\n  print(\"Inside:\", y)\n  return x + y\n\nfor y in range(3):\n  print(\"Result:\", impure_func(y))\n```\n\nExample:\n```text\nInside: 0\nResult: 0\nInside: 1\nResult: 2\nInside: 2\nResult: 4\n```\n\nExample:\n```text\nInside: 0\nResult: 0\nResult: 1\nResult: 2\n```\n\nExample:\n```text\n>>> from jax import jit\n>>> import jax.numpy as jnp\n>>> def f(x):\n...   return jnp.log(jnp.sqrt(x))\n>>> x = jnp.pi\n>>> print(f(x))\n0.572365\n```\n\nExample:\n```text\n>>> print(jit(f)(x))\n0.5723649\n```\n\nExample:\n```text\n>>> def f(x):\n...   return jnp.exp(x) / jnp.exp(x)\n>>> x = 100.0\n>>> print(f(x))\nnan\n```\n\nExample:\n```text\n>>> print(jit(f)(x))\n1.0\n```\n\nExample:\n```text\ndef my_log(x):\n  return jnp.where(x > 0., jnp.log(x), 0.)\n\nmy_log(0.) ==> 0.  # Ok\njax.grad(my_log)(0.)  ==> NaN\n```\n\nExample:\n```text\ndef safe_for_grad_log(x):\n  return jnp.log(jnp.where(x > 0., x, 1.))\n\nsafe_for_grad_log(0.) ==> 0.  # Ok\njax.grad(safe_for_grad_log)(0.)  ==> 0.  # Ok\n```\n\nExample:\n```text\ndef my_log_or_y(x, y):\n  \"\"\"Return log(x) if x > 0 or y\"\"\"\n  return jnp.where(x > 0., jnp.log(jnp.where(x > 0., x, 1.)), y)\n```\n\nExample:\n```text\nimport jax\nimport numpy as np\nimport jax.numpy as jnp\n\ndef f(x):\n  return (x > 0).astype(float)\n\ndf = jax.vmap(jax.grad(f))\n\nx = jnp.array([-1.0, -0.5, 0.0, 0.5, 1.0])\n\nprint(f\"f(x)  = {f(x)}\")\n# f(x)  = [0. 0. 0. 1. 1.]\n\nprint(f\"df(x) = {df(x)}\")\n# df(x) = [0. 0. 0. 0. 0.]\n```\n\nExample:\n```text\ndef g(x):\n  return jax.nn.sigmoid(x)\n\ndg = jax.vmap(jax.grad(g))\n\nx = jnp.array([-10.0, -1.0, 0.0, 1.0, 10.0])\n\nwith np.printoptions(suppress=True, precision=2):\n  print(f\"g(x)  = {g(x)}\")\n  # g(x)  = [0.   0.27 0.5  0.73 1.  ]\n\n  print(f\"dg(x) = {dg(x)}\")\n  # dg(x) = [0.   0.2  0.25 0.2  0.  ]\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  print(type(x))\n  return x\n\nf(jnp.arange(5))\n```\n\nExample:\n```text\n<class 'jax.interpreters.partial_eval.DynamicJaxprTracer'>\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.860Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":14,"totalLines":141,"estimatedTokens":486}}132{"id":"doc-type_promotion_semantics_jax_documentation-e749e72b","source":"documentation","title":"Type promotion semantics — JAX documentation","url":"https://docs.jax.dev/en/latest/type_promotion.html","text":"Example:\n```text\n>>> x = jnp.arange(5, dtype='int8')\n>>> 2 * x\nArray([0, 2, 4, 6, 8], dtype=int8)\n```\n\nExample:\n```text\n>>> jnp.int32(2) * x\nArray([0, 2, 4, 6, 8], dtype=int32)\n```\n\nExample:\n```text\n>>> jnp.asarray(2)\nArray(2, dtype=int32, weak_type=True)\n```\n\nExample:\n```text\n>>> jnp.asarray(2, dtype='int32')\nArray(2, dtype=int32)\n```\n\nExample:\n```text\n>>> x = jnp.float32(1)\n>>> y = jnp.int32(1)\n>>> with jax.numpy_dtype_promotion('strict'):\n...   z = x + y  \n...\nTraceback (most recent call last):\nTypePromotionError: Input dtypes ('float32', 'int32') have no available implicit\ndtype promotion path when jax_numpy_dtype_promotion=strict. Try explicitly casting\ninputs to the desired output type, or set jax_numpy_dtype_promotion=standard.\n```\n\nExample:\n```text\n>>> with jax.numpy_dtype_promotion('strict'):\n...   z = x + 1\n>>> print(z)\n2.0\n```\n\nExample:\n```text\njax.config.update('jax_numpy_dtype_promotion', 'strict')\n```\n\nExample:\n```text\njax.config.update('jax_numpy_dtype_promotion', 'standard')\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.861Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":8,"totalLines":57,"estimatedTokens":256}}133{"id":"doc-writing_custom_jaxpr_interpreters_in_jax_jax_doc-7ba72626","source":"documentation","title":"Writing custom Jaxpr interpreters in JAX — JAX documentation","url":"https://docs.jax.dev/en/latest/notebooks/Writing_custom_interpreters_in_Jax.html","text":"Example:\n```text\nimport jax\nimport jax.numpy as jnp\nfrom jax import jit, grad, vmap\nfrom jax import random\n```\n\nExample:\n```text\nx = random.normal(random.key(0), (5000, 5000))\ndef f(w, b, x):\n  return jnp.tanh(jnp.dot(x, w) + b)\nfast_f = jit(f)\n```\n\nExample:\n```text\ndef examine_jaxpr(closed_jaxpr):\n  jaxpr = closed_jaxpr.jaxpr\n  print(\"invars:\", jaxpr.invars)\n  print(\"outvars:\", jaxpr.outvars)\n  print(\"constvars:\", jaxpr.constvars)\n  for eqn in jaxpr.eqns:\n    print(\"equation:\", eqn.invars, eqn.primitive, eqn.outvars, eqn.params)\n  print()\n  print(\"jaxpr:\", jaxpr)\n\ndef foo(x):\n  return x + 1\nprint(\"foo\")\nprint(\"=====\")\nexamine_jaxpr(jax.make_jaxpr(foo)(5))\n\nprint()\n\ndef bar(w, b, x):\n  return jnp.dot(w, x) + b + jnp.ones(5), x\nprint(\"bar\")\nprint(\"=====\")\nexamine_jaxpr(jax.make_jaxpr(bar)(jnp.ones((5, 10)), jnp.ones(5), jnp.ones(10)))\n```\n\nExample:\n```text\nfoo\n=====\ninvars: [Var(id=131066573615600):int32[]]\noutvars: [Var(id=131066573618384):int32[]]\nconstvars: []\nequation: [Var(id=131066573615600):int32[], Literal(TypedInt(1, dtype=int32))] add [Var(id=131066573618384):int32[]] {}\n\njaxpr: { lambda ; a:i32[]. let b:i32[] = add a 1:i32[] in (b,) }\n\nbar\n=====\ninvars: [Var(id=131066573905712):float32[5,10], Var(id=131066573905904):float32[5], Var(id=131066573906096):float32[10]]\noutvars: [Var(id=131066573910512):float32[5], Var(id=131066573906096):float32[10]]\nconstvars: []\nequation: [Var(id=131066573905712):float32[5,10], Var(id=131066573906096):float32[10]] dot_general [Var(id=131066573906960):float32[5]] {'dimension_numbers': (((1,), (0,)), ((), ())), 'precision': None, 'preferred_element_type': dtype('float32'), 'out_sharding': None}\nequation: [Var(id=131066573906960):float32[5], Var(id=131066573905904):float32[5]] add [Var(id=131066573908400):float32[5]] {}\nequation: [Literal(1.0)] broadcast_in_dim [Var(id=131066573909408):float32[5]] {'shape': (5,), 'broadcast_dimensions': (), 'sharding': None}\nequation: [Var(id=131066573908400):float32[5], Var(id=131066573909408):float32[5]] add [Var(id=131066573910512):float32[5]] {}\n\njaxpr: { lambda ; a:f32[5,10] b:f32[5] c:f32[10]. let\n    d:f32[5] = dot_general[\n      dimension_numbers=(([1], [0]), ([], []))\n      preferred_element_type=float32\n    ] a c\n    e:f32[5] = add d b\n    f:f32[5] = broadcast_in_dim 1.0:f32[]\n    g:f32[5] = add e f\n  in (g, c) }\n```\n\nExample:\n```text\ndef f(x):\n  return jnp.exp(jnp.tanh(x))\nf_inv = inverse(f)\nassert jnp.allclose(f_inv(f(1.0)), 1.0)\n```\n\nExample:\n```text\n# Importing Jax functions useful for tracing/interpreting.\nfrom functools import wraps\n\nfrom jax import lax\nfrom jax.extend import core\nfrom jax._src.util import safe_map\n```\n\nExample:\n```text\ndef f(x):\n  return jnp.exp(jnp.tanh(x))\n\nclosed_jaxpr = jax.make_jaxpr(f)(jnp.ones(5))\nprint(closed_jaxpr.jaxpr)\nprint(closed_jaxpr.literals)\n```\n\nExample:\n```text\n{ lambda ; a:f32[5]. let b:f32[5] = tanh a; c:f32[5] = exp b in (c,) }\n[]\n```\n\nExample:\n```text\ndef eval_jaxpr(jaxpr, consts, *args):\n  # Mapping from variable -> value\n  env = {}\n\n  def read(var):\n    # Literals are values baked into the Jaxpr\n    if type(var) is core.Literal:\n      return var.val\n    return env[var]\n\n  def write(var, val):\n    env[var] = val\n\n  # Bind args and consts to environment\n  safe_map(write, jaxpr.invars, args)\n  safe_map(write, jaxpr.constvars, consts)\n\n  # Loop through equations and evaluate primitives using `bind`\n  for eqn in jaxpr.eqns:\n    # Read inputs to equation from environment\n    invals = safe_map(read, eqn.invars)\n    # `bind` is how a primitive is called\n    outvals = eqn.primitive.bind(*invals, **eqn.params)\n    # Primitives may return multiple outputs or not\n    if not eqn.primitive.multiple_results:\n      outvals = [outvals]\n    # Write the results of the primitive into the environment\n    safe_map(write, eqn.outvars, outvals)\n  # Read the final result of the Jaxpr from the environment\n  return safe_map(read, jaxpr.outvars)\n```\n\nExample:\n```text\nclosed_jaxpr = jax.make_jaxpr(f)(jnp.ones(5))\neval_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.literals, jnp.ones(5))\n```\n\nExample:\n```text\n[Array([2.1416876, 2.1416876, 2.1416876, 2.1416876, 2.1416876], dtype=float32)]\n```\n\nExample:\n```text\ninverse_registry = {}\n```\n\nExample:\n```text\ninverse_registry[lax.exp_p] = jnp.log\ninverse_registry[lax.tanh_p] = jnp.arctanh\n```\n\nExample:\n```text\ndef inverse(fun):\n  @wraps(fun)\n  def wrapped(*args, **kwargs):\n    # Since we assume unary functions, we won't worry about flattening and\n    # unflattening arguments.\n    closed_jaxpr = jax.make_jaxpr(fun)(*args, **kwargs)\n    out = inverse_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.literals, *args)\n    return out[0]\n  return wrapped\n```\n\nExample:\n```text\ndef inverse_jaxpr(jaxpr, consts, *args):\n  env = {}\n\n  def read(var):\n    if type(var) is core.Literal:\n      return var.val\n    return env[var]\n\n  def write(var, val):\n    env[var] = val\n  # Args now correspond to Jaxpr outvars\n  safe_map(write, jaxpr.outvars, args)\n  safe_map(write, jaxpr.constvars, consts)\n\n  # Looping backward\n  for eqn in jaxpr.eqns[::-1]:\n    #  outvars are now invars\n    invals = safe_map(read, eqn.outvars)\n    if eqn.primitive not in inverse_registry:\n      raise NotImplementedError(\n          f\"{eqn.primitive} does not have registered inverse.\")\n    # Assuming a unary function\n    outval = inverse_registry[eqn.primitive](*invals)\n    safe_map(write, eqn.invars, [outval])\n  return safe_map(read, jaxpr.invars)\n```\n\nExample:\n```text\njax.make_jaxpr(inverse(f))(f(1.))\n```\n\nExample:\n```text\n{ lambda ; a:f32[]. let b:f32[] = log a; c:f32[] = atanh b in (c,) }\n```\n\nExample:\n```text\njit(vmap(grad(inverse(f))))((jnp.arange(5) + 1.) / 5.)\n```\n\nExample:\n```text\nArray([-3.1440797, 15.584931 ,  2.2551253,  1.3155028,  1.       ],      dtype=float32, weak_type=True)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.861Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":19,"totalLines":228,"estimatedTokens":1448}}134{"id":"doc-design_of_type_promotion_semantics_for_jax_jax_d-207b1b31","source":"documentation","title":"Design of Type Promotion Semantics for JAX — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/9407-type-promotion.html","text":"Example:\n```text\nimport numpy as np\nnp.dtype(np.int32(1) + np.float32(1))\n```\n\nExample:\n```text\ndtype('float64')\n```\n\nExample:\n```text\nimport pandas as pd\ntypes = [int, float, complex]\nname = lambda t: t.__name__\npd.DataFrame([[name(type(t1(1) + t2(1))) for t1 in types] for t2 in types],\n             index=[name(t) for t in types], columns=[name(t) for t in types])\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {'int': ['float'], 'float': ['complex']}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {'int': [0, 0], 'float': [1, 0], 'complex': [2, 0]}\nfig, ax = plt.subplots(figsize=(8, 2))\nnx.draw(graph, with_labels=True, node_size=4000, node_color='lightgray', pos=pos, ax=ax, arrowsize=20)\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\n\nfig, ax = plt.subplots(1, 2, figsize=(10, 2))\n\nlattice = {'A': ['B', 'C']}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {'A': [0, 0], 'B': [1, 0.5], 'C': [1, -0.5]}\nnx.draw(graph, with_labels=True, node_size=2000, node_color='lightgray', pos=pos, ax=ax[0], arrowsize=20)\nax[0].set(xlim=[-0.5, 1.5], ylim=[-1, 1])\n\nlattice = {'A': ['C', 'D'], 'B': ['C', 'D']}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {'A': [0, 0.5], 'B': [0, -0.5], 'C': [1, 0.5], 'D': [1, -0.5]}\nnx.draw(graph, with_labels=True, node_size=2000, node_color='lightgray', pos=pos, ax=ax[1], arrowsize=20)\nax[1].set(xlim=[-0.5, 1.5], ylim=[-1, 1]);\n```\n\nExample:\n```text\nimport numpy as np\na, b, c = np.int8(1), np.uint8(1), np.float16(1)\nprint(np.dtype((a + b) + c))\nprint(np.dtype(a + (b + c)))\n```\n\nExample:\n```text\nfloat32\nfloat16\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'u8': ['u16'], 'u16': ['u32'], 'u32': ['u64'],\n  'i8': ['i16'], 'i16': ['i32'], 'i32': ['i64'],\n  'f16': ['f32'], 'f32': ['f64'],\n  'c64': ['c128']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'u8': [0, 0], 'u16': [1, 0], 'u32': [2, 0], 'u64': [3, 0],\n  'i8': [0, 1], 'i16': [1, 1], 'i32': [2, 1], 'i64': [3, 1],\n  'f16': [1, 2], 'f32': [2, 2], 'f64': [3, 2],\n  'c64': [2, 3], 'c128': [3, 3],\n}\nfig, ax = plt.subplots(figsize=(6, 4))\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\n```\n\nExample:\n```text\nx = np.int8(0)  # int8 scalar\ny = 1  # Python int = int64 scalar\n(x + y).dtype\n```\n\nExample:\n```text\ndtype('int64')\n```\n\nExample:\n```text\nx = np.zeros(1, dtype='int8')  # int8 array\ny = 1  # Python int = int64 scalar\n(x + y).dtype\n```\n\nExample:\n```text\ndtype('int8')\n```\n\nExample:\n```text\nx = np.zeros(1, dtype='int8')  # int8 array\ny = 1000  # int64 scalar\n(x + y).dtype\n```\n\nExample:\n```text\ndtype('int16')\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'i8*': ['i16*'], 'i16*': ['i32*'], 'i32*': ['i64*'], 'i64*': ['i8'],\n  'i8': ['i16'], 'i16': ['i32'], 'i32': ['i64']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'i8*': [0, 1], 'i16*': [2, 1], 'i32*': [4, 1], 'i64*': [6, 1],\n  'i8': [9, 1], 'i16': [11, 1], 'i32': [13, 1], 'i64': [15, 1],\n}\nfig, ax = plt.subplots(figsize=(12, 4))\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\nax.text(3, 1.6, \"Scalar Types\", ha='center', fontsize=14)\nax.text(12, 1.6, \"Array Types\", ha='center', fontsize=14)\nax.set_ylim(-1, 3);\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'u*': ['u8'], 'u8': ['u16'], 'u16': ['u32'], 'u32': ['u64'],\n  'i*': ['i8'], 'i8': ['i16'], 'i16': ['i32'], 'i32': ['i64'],\n  'f*': ['f16'], 'f16': ['f32'], 'f32': ['f64'],\n  'c*': ['c64'], 'c64': ['c128']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'u*': [0, 0], 'u8': [3, 0], 'u16': [5, 0], 'u32': [7, 0], 'u64': [9, 0],\n  'i*': [0, 1], 'i8': [3, 1], 'i16': [5, 1], 'i32': [7, 1], 'i64': [9, 1],\n  'f*': [0, 2], 'f16': [5, 2], 'f32': [7, 2], 'f64': [9, 2],\n  'c*': [0, 3], 'c64': [7, 3], 'c128': [9, 3],\n}\nfig, ax = plt.subplots(figsize=(6, 4))\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\n```\n\nExample:\n```text\nfor dtype in [np.int8, np.int16, np.int32, np.int64]:\n  x = np.arange(10, dtype=dtype)\n  assert (x + 2).dtype == dtype\n```\n\nExample:\n```text\n3 * (x + 1) ** 2\n```\n\nExample:\n```text\nnp.int32(3) * (x + np.int32(1)) ** np.int32(2)\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'i*': ['f*', 'u8', 'i8'], 'f*': ['c*', 'f16'], 'c*': ['c64'],\n  'u8': ['u16'], 'u16': ['u32'], 'u32': ['u64'],\n  'i8': ['i16'], 'i16': ['i32'], 'i32': ['i64'],\n  'f16': ['f32'], 'f32': ['f64'],\n  'c64': ['c128']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'i*': [-1.25, 0.5], 'f*': [-0.5, 2], 'c*': [0, 3],\n  'u8': [0.5, 0], 'u16': [1.5, 0], 'u32': [2.5, 0], 'u64': [3.5, 0],\n  'i8': [0, 1], 'i16': [1, 1], 'i32': [2, 1], 'i64': [3, 1],\n  'f16': [0.5, 2], 'f32': [1.5, 2], 'f64': [2.5, 2],\n  'c64': [2, 3], 'c128': [3, 3],\n}\nfig, ax = plt.subplots(figsize=(6, 5))\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'i*': ['f*', 'u8', 'i8'], 'f*': ['c*', 'f16'], 'c*': ['c64'],\n  'u8': ['u16'], 'u16': ['u32'], 'u32': ['u64'],\n  'i8': ['i16'], 'i16': ['i32'], 'i32': ['i64'],\n  'f16': ['f32'], 'f32': ['f64', 'c64'], 'f64': ['c128'],\n  'c64': ['c128']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'i*': [-1.25, 0.5], 'f*': [-0.5, 2], 'c*': [0, 3],\n  'u8': [0.5, 0], 'u16': [1.5, 0], 'u32': [2.5, 0], 'u64': [3.5, 0],\n  'i8': [0, 1], 'i16': [1, 1], 'i32': [2, 1], 'i64': [3, 1],\n  'f16': [0.5, 2], 'f32': [1.5, 2], 'f64': [2.5, 2],\n  'c64': [2, 3], 'c128': [3, 3],\n}\nfig, ax = plt.subplots(figsize=(6, 5))\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'i*': ['f*', 'u8', 'i8'], 'f*': ['c*', 'f16'], 'c*': ['c64'],\n  'u8': ['u16', 'i16'], 'u16': ['u32', 'i32'], 'u32': ['u64', 'i64'],\n  'i8': ['i16'], 'i16': ['i32'], 'i32': ['i64'],\n  'f16': ['f32'], 'f32': ['f64', 'c64'], 'f64': ['c128'],\n  'c64': ['c128']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'i*': [-1.25, 0.5], 'f*': [-0.5, 2], 'c*': [0, 3],\n  'u8': [0.5, 0], 'u16': [1.5, 0], 'u32': [2.5, 0], 'u64': [3.5, 0],\n  'i8': [0, 1], 'i16': [1, 1], 'i32': [2, 1], 'i64': [3, 1],\n  'f16': [0.5, 2], 'f32': [1.5, 2], 'f64': [2.5, 2],\n  'c64': [2, 3], 'c128': [3, 3],\n}\nfig, ax = plt.subplots(figsize=(6, 5))\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\n```\n\nExample:\n```text\n(np.uint64(1) + np.int64(1)).dtype\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'i*': ['f*', 'u8', 'i8'], 'f*': ['c*', 'f16'], 'c*': ['c64'],\n  'u8': ['u16', 'i16', 'f16'], 'u16': ['u32', 'i32', 'f32'], 'u32': ['u64', 'i64', 'f64'],\n  'i8': ['i16', 'f16'], 'i16': ['i32', 'f32'], 'i32': ['i64', 'f64'],\n  'f16': ['f32'], 'f32': ['f64', 'c64'], 'f64': ['c128'],\n  'c64': ['c128']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'i*': [-1.25, 0.5], 'f*': [-0.5, 2], 'c*': [0, 3],\n  'u8': [0.5, 0], 'u16': [1.5, 0], 'u32': [2.5, 0], 'u64': [3.5, 0],\n  'i8': [0, 1], 'i16': [1, 1], 'i32': [2, 1], 'i64': [3, 1],\n  'f16': [0.5, 2], 'f32': [1.5, 2], 'f64': [2.5, 2],\n  'c64': [2, 3], 'c128': [3, 3],\n}\nfig, ax = plt.subplots(figsize=(6, 5))\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'i*': ['f*', 'u8', 'i8'], 'f*': ['c*', 'f16'], 'c*': ['c64'],\n  'u8': ['u16', 'i16'], 'u16': ['u32', 'i32'], 'u32': ['u64', 'i64'],\n  'i8': ['i16', 'f16'], 'i16': ['i32', 'f32'], 'i32': ['i64', 'f64'],\n  'f16': ['f32'], 'f32': ['f64', 'c64'], 'f64': ['c128'],\n  'c64': ['c128']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'i*': [-1.25, 0.5], 'f*': [-0.5, 2], 'c*': [0, 3],\n  'u8': [0.5, 0], 'u16': [1.5, 0], 'u32': [2.5, 0], 'u64': [3.5, 0],\n  'i8': [0, 1], 'i16': [1, 1], 'i32': [2, 1], 'i64': [3, 1],\n  'f16': [0.5, 2], 'f32': [1.5, 2], 'f64': [2.5, 2],\n  'c64': [2, 3], 'c128': [3, 3],\n}\nfig, ax = plt.subplots(figsize=(6, 5))\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'i*': ['f*', 'u8', 'i8'], 'f*': ['c*', 'f16'], 'c*': ['c64'],\n  'u8': ['u16', 'i16'], 'u16': ['u32', 'i32'], 'u32': ['u64', 'i64'],\n  'i8': ['i16'], 'i16': ['f16', 'i32'], 'i32': ['f32', 'i64'], 'i64': ['f64'],\n  'f16': ['f32'], 'f32': ['f64', 'c64'], 'f64': ['c128'],\n  'c64': ['c128']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'i*': [-1.25, 0.5], 'f*': [-0.5, 2], 'c*': [0, 3],\n  'u8': [0.5, 0], 'u16': [1.5, 0], 'u32': [2.5, 0], 'u64': [3.5, 0],\n  'i8': [0, 1], 'i16': [1, 1], 'i32': [2, 1], 'i64': [3, 1],\n  'f16': [1.5, 2], 'f32': [2.5, 2], 'f64': [3.5, 2],\n  'c64': [3, 3], 'c128': [4, 3],\n}\nfig, ax = plt.subplots(figsize=(6, 5))\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'i*': ['u8', 'i8'], 'f*': ['c*', 'f16'], 'c*': ['c64'],\n  'u8': ['u16', 'i16'], 'u16': ['u32', 'i32'], 'u32': ['u64', 'i64'],\n  'i8': ['i16'], 'i16': ['i32'], 'i32': ['i64'], 'i64': ['f*'],\n  'f16': ['f32'], 'f32': ['f64', 'c64'], 'f64': ['c128'],\n  'c64': ['c128']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'i*': [-1.25, 0.5], 'f*': [-0.5, 2], 'c*': [0, 3],\n  'u8': [0.5, 0], 'u16': [1.5, 0], 'u32': [2.5, 0], 'u64': [3.5, 0],\n  'i8': [0, 1], 'i16': [1, 1], 'i32': [2, 1], 'i64': [3, 1],\n  'f16': [1.5, 2], 'f32': [2.5, 2], 'f64': [3.5, 2],\n  'c64': [3, 3], 'c128': [4, 3],\n}\nfig, ax = plt.subplots(figsize=(6, 5))\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'i*': ['u8', 'i8'], 'f*': ['c*', 'f16', 'bf16'], 'c*': ['c64'],\n  'u8': ['u16', 'i16'], 'u16': ['u32', 'i32'], 'u32': ['u64', 'i64'],\n  'i8': ['i16'], 'i16': ['i32'], 'i32': ['i64'], 'i64': ['f*'],\n  'f16': ['f32'], 'bf16': ['f32'], 'f32': ['f64', 'c64'], 'f64': ['c128'],\n  'c64': ['c128']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'i*': [-1.25, 0.5], 'f*': [-0.5, 2], 'c*': [0, 3],\n  'u8': [0.5, 0], 'u16': [1.5, 0], 'u32': [2.5, 0], 'u64': [3.5, 0],\n  'i8': [0, 1], 'i16': [1, 1], 'i32': [2, 1], 'i64': [3, 1],\n  'f16': [1.8, 1.7], 'bf16': [1.8, 2.3], 'f32': [3.0, 2], 'f64': [4.0, 2],\n  'c64': [3.5, 3], 'c128': [4.5, 3],\n}\nfig, ax = plt.subplots(figsize=(6, 5))\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\n```\n\nExample:\n```text\n#@title\nimport networkx as nx\nimport matplotlib.pyplot as plt\nlattice = {\n  'i*': ['u8', 'i8'], 'f*': ['c*', 'f16', 'bf16'], 'c*': ['c64'],\n  'u8': ['u16', 'i16'], 'u16': ['u32', 'i32'], 'u32': ['u64', 'i64'], 'u64': ['f*'],\n  'i8': ['i16'], 'i16': ['i32'], 'i32': ['i64'], 'i64': ['f*'],\n  'f16': ['f32'], 'bf16': ['f32'], 'f32': ['f64', 'c64'], 'f64': ['c128'],\n  'c64': ['c128']\n}\ngraph = nx.from_dict_of_lists(lattice, create_using=nx.DiGraph)\npos = {\n  'i*': [-1.25, 0.5], 'f*': [4.5, 0.5], 'c*': [5, 1.5],\n  'u8': [0.5, 0], 'u16': [1.5, 0], 'u32': [2.5, 0], 'u64': [3.5, 0],\n  'i8': [0, 1], 'i16': [1, 1], 'i32': [2, 1], 'i64': [3, 1],\n  'f16': [5.75, 0.8], 'bf16': [5.75, 0.2], 'f32': [7, 0.5], 'f64': [8, 0.5],\n  'c64': [7.5, 1.5], 'c128': [8.5, 1.5],\n}\nfig, ax = plt.subplots(figsize=(10, 4))\nax.set_ylim(-0.5, 2)\nnx.draw(graph, with_labels=True, node_size=1500, node_color='lightgray', pos=pos, ax=ax)\n# ax.patches[12].set_linestyle((0, (2, 4)))\n```\n\nExample:\n```text\n# @title\n\nimport numpy as np\nimport pandas as pd\nfrom IPython import display\n\nnp_dtypes = {\n  'b': np.bool_,\n  'u8': np.uint8, 'u16': np.uint16, 'u32': np.uint32, 'u64': np.uint64,\n  'i8': np.int8, 'i16': np.int16, 'i32': np.int32, 'i64': np.int64,\n  'bf16': 'invalid', 'f16': np.float16, 'f32': np.float32, 'f64': np.float64,\n  'c64': np.complex64, 'c128': np.complex128,\n  'i*': int, 'f*': float, 'c*': complex}\n\nnp_dtype_to_code = {val: key for key, val in np_dtypes.items()}\n\ndef make_np_zero(dtype):\n  if dtype in {int, float, complex}:\n    return dtype(0)\n  else:\n    return np.zeros(1, dtype=dtype)\n\ndef np_result_code(dtype1, dtype2):\n  try:\n    out = np.add(make_np_zero(dtype1), make_np_zero(dtype2))\n  except TypeError:\n    return '-'\n  else:\n    if type(out) in {int, float, complex}:\n      return np_dtype_to_code[type(out)]\n    else:\n      return np_dtype_to_code[out.dtype.type]\n\n\ngrid = [[np_result_code(dtype1, dtype2)\n         for dtype2 in np_dtypes.values()]\n        for dtype1 in np_dtypes.values()]\ntable = pd.DataFrame(grid, index=np_dtypes.keys(), columns=np_dtypes.keys())\ndisplay.HTML(table.to_html())\n```\n\nExample:\n```text\n# @title\n\nimport tensorflow as tf\nimport pandas as pd\nfrom IPython import display\n\ntf_dtypes = {\n  'b': tf.bool,\n  'u8': tf.uint8, 'u16': tf.uint16, 'u32': tf.uint32, 'u64': tf.uint64,\n  'i8': tf.int8, 'i16': tf.int16, 'i32': tf.int32, 'i64': tf.int64,\n  'bf16': tf.bfloat16, 'f16': tf.float16, 'f32': tf.float32, 'f64': tf.float64,\n  'c64': tf.complex64, 'c128': tf.complex128,\n  'i*': int, 'f*': float, 'c*': complex}\n\ntf_dtype_to_code = {val: key for key, val in tf_dtypes.items()}\n\ndef make_tf_zero(dtype):\n  if dtype in {int, float, complex}:\n    return dtype(0)\n  else:\n    return tf.zeros(1, dtype=dtype)\n\ndef result_code(dtype1, dtype2):\n  try:\n    out = tf.add(make_tf_zero(dtype1), make_tf_zero(dtype2))\n  except (TypeError, tf.errors.InvalidArgumentError):\n    return '-'\n  else:\n    if type(out) in {int, float, complex}:\n      return tf_dtype_to_code[type(out)]\n    else:\n      return tf_dtype_to_code[out.dtype]\n\n\ngrid = [[result_code(dtype1, dtype2)\n         for dtype2 in tf_dtypes.values()]\n        for dtype1 in tf_dtypes.values()]\ntable = pd.DataFrame(grid, index=tf_dtypes.keys(), columns=tf_dtypes.keys())\ndisplay.HTML(table.to_html())\n```\n\nExample:\n```text\n# @title\nimport torch\nimport pandas as pd\nfrom IPython import display\n\ntorch_dtypes = {\n  'b': torch.bool,\n  'u8': torch.uint8, 'u16': 'invalid', 'u32': 'invalid', 'u64': 'invalid',\n  'i8': torch.int8, 'i16': torch.int16, 'i32': torch.int32, 'i64': torch.int64,\n  'bf16': torch.bfloat16, 'f16': torch.float16, 'f32': torch.float32, 'f64': torch.float64,\n  'c64': torch.complex64, 'c128': torch.complex128,\n  'i*': int, 'f*': float, 'c*': complex}\n\ntorch_dtype_to_code = {val: key for key, val in torch_dtypes.items()}\n\ndef make_torch_zero(dtype):\n  if dtype in {int, float, complex}:\n    return dtype(0)\n  else:\n    return torch.zeros(1, dtype=dtype)\n\ndef torch_result_code(dtype1, dtype2):\n  try:\n    out = torch.add(make_torch_zero(dtype1), make_torch_zero(dtype2))\n  except TypeError:\n    return '-'\n  else:\n    if type(out) in {int, float, complex}:\n      return torch_dtype_to_code[type(out)]\n    else:\n      return torch_dtype_to_code[out.dtype]\n\n\ngrid = [[torch_result_code(dtype1, dtype2)\n         for dtype2 in torch_dtypes.values()]\n        for dtype1 in torch_dtypes.values()]\ntable = pd.DataFrame(grid, index=torch_dtypes.keys(), columns=torch_dtypes.keys())\ndisplay.HTML(table.to_html())\n```\n\nExample:\n```text\n# @title\nimport jax\nimport jax.numpy as jnp\nimport pandas as pd\nfrom IPython import display\njax.config.update('jax_enable_x64', True)\n\njnp_dtypes = {\n  'b': jnp.bool_.dtype,\n  'u8': jnp.uint8.dtype, 'u16': jnp.uint16.dtype, 'u32': jnp.uint32.dtype, 'u64': jnp.uint64.dtype,\n  'i8': jnp.int8.dtype, 'i16': jnp.int16.dtype, 'i32': jnp.int32.dtype, 'i64': jnp.int64.dtype,\n  'bf16': jnp.bfloat16.dtype, 'f16': jnp.float16.dtype, 'f32': jnp.float32.dtype, 'f64': jnp.float64.dtype,\n  'c64': jnp.complex64.dtype, 'c128': jnp.complex128.dtype,\n  'i*': int, 'f*': float, 'c*': complex}\n\n\njnp_dtype_to_code = {val: key for key, val in jnp_dtypes.items()}\n\ndef make_jnp_zero(dtype):\n  if dtype in {int, float, complex}:\n    return dtype(0)\n  else:\n    return jnp.zeros((), dtype=dtype)\n\ndef jnp_result_code(dtype1, dtype2):\n  try:\n    out = jnp.add(make_jnp_zero(dtype1), make_jnp_zero(dtype2))\n  except TypeError:\n    return '-'\n  else:\n    if hasattr(out, 'aval') and out.aval.weak_type:\n      return out.dtype.kind + '*'\n    elif type(out) in {int, float, complex}:\n      return jnp_dtype_to_code[type(out)]\n    else:\n      return jnp_dtype_to_code[out.dtype]\n\ngrid = [[jnp_result_code(dtype1, dtype2)\n         for dtype2 in jnp_dtypes.values()]\n        for dtype1 in jnp_dtypes.values()]\ntable = pd.DataFrame(grid, index=jnp_dtypes.keys(), columns=jnp_dtypes.keys())\ndisplay.HTML(table.to_html())\n```\n\nExample:\n```text\n# @title\nimport jax\nimport jax.numpy as jnp\nimport pandas as pd\nfrom IPython import display\njax.config.update('jax_enable_x64', True)\n\njnp_dtypes = {\n  'b': jnp.bool_.dtype,\n  'u8': jnp.uint8.dtype, 'u16': jnp.uint16.dtype, 'u32': jnp.uint32.dtype, 'u64': jnp.uint64.dtype,\n  'i8': jnp.int8.dtype, 'i16': jnp.int16.dtype, 'i32': jnp.int32.dtype, 'i64': jnp.int64.dtype,\n  'bf16': jnp.bfloat16.dtype, 'f16': jnp.float16.dtype, 'f32': jnp.float32.dtype, 'f64': jnp.float64.dtype,\n  'c64': jnp.complex64.dtype, 'c128': jnp.complex128.dtype,\n  'i*': int, 'f*': float, 'c*': complex}\n\n\njnp_dtype_to_code = {val: key for key, val in jnp_dtypes.items()}\n\ndef make_jnp_zero(dtype):\n  if dtype in {int, float, complex}:\n    return dtype(0)\n  else:\n    return jnp.zeros((), dtype=dtype)\n\ndef jnp_result_code(dtype1, dtype2):\n  try:\n    out = jax.lax.add(make_jnp_zero(dtype1), make_jnp_zero(dtype2))\n  except TypeError:\n    return '-'\n  else:\n    if hasattr(out, 'aval') and out.aval.weak_type:\n      return out.dtype.kind + '*'\n    elif type(out) in {int, float, complex}:\n      return jnp_dtype_to_code[type(out)]\n    else:\n      return jnp_dtype_to_code[out.dtype]\n\ngrid = [[jnp_result_code(dtype1, dtype2)\n         for dtype2 in jnp_dtypes.values()]\n        for dtype1 in jnp_dtypes.values()]\ntable = pd.DataFrame(grid, index=jnp_dtypes.keys(), columns=jnp_dtypes.keys())\ndisplay.HTML(table.to_html())\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.864Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":34,"totalLines":629,"estimatedTokens":4661}}135{"id":"doc-configuration_options_jax_documentation-e1b78901","source":"documentation","title":"Configuration Options — JAX documentation","url":"https://docs.jax.dev/en/latest/config_options.html","text":"Example:\n```text\nexport JAX_ENABLE_X64=True\npython my_program.py\n```\n\nExample:\n```text\nimport jax\njax.config.update(\"jax_enable_x64\", True)\n```\n\nExample:\n```text\n# In your code:\nimport jax\njax.config.parse_flags_with_absl()\n```\n\nExample:\n```text\n# When running:\npython my_program.py --jax_enable_x64=True\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.870Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":26,"estimatedTokens":81}}136{"id":"doc-change_log_jax_documentation-af38ec8f","source":"documentation","title":"Change log — JAX documentation","url":"https://docs.jax.dev/en/latest/changelog.html","text":"Example:\n```text\n@functools.partial(jax.jit, static_argnames=['n'])\ndef f(x, n):\n  ...\n```\n\nExample:\n```text\n@jax.jit(static_argnames=['n'])\ndef f(x, n):\n  ...\n```\n\nExample:\n```text\n@jax.jit\ndef f(x):\n  return x\n\n# inp1.sharding is of type SingleDeviceSharding\ninp1 = jnp.arange(8)\nf(inp1)\n\nmesh = jax.make_mesh((1,), ('x',))\n# inp2.sharding is of type NamedSharding\ninp2 = jax.device_put(jnp.arange(8), NamedSharding(mesh, P('x')))\nf(inp2)  # tracing cache miss\n```\n\nExample:\n```text\n@jtu.with_config(jax_numpy_rank_promotion='allow')\nclass MyTestCase(jtu.JaxTestCase):\n  ...\n```\n\nExample:\n```text\npip install --upgrade pip\n\n# Installs the wheel compatible with CUDA 11 and cuDNN 8.2 or newer.\npip install --upgrade \"jax[cuda]\" -f https://storage.googleapis.com/jax-releases/jax_releases.html\n\n# Installs the wheel compatible with Cuda 11 and cudnn 8.2 or newer.\npip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_releases.html\n\n# Installs the wheel compatible with Cuda 11 and cudnn 8.0.5 or newer.\npip install jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_releases.html\n```\n\nExample:\n```text\nkey = random.PRNGKey(-1).at[0].set(0xFFFFFFFF)\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.878Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":6,"totalLines":57,"estimatedTokens":305}}137{"id":"doc-jax_remat_jax_checkpoint_changes_what_you_need_t-77ff98bd","source":"documentation","title":"jax.remat / jax.checkpoint changes: what you need to know — JAX documentation","url":"https://docs.jax.dev/en/latest/jep/11830-new-remat-checkpoint.html","text":"Example:\n```text\nfrom functools import partial\nimport jax\n\ndef apply_layer(W, x):\n  return jnp.sin(jnp.dot(W, x))\n\n@partial(jax.checkpoint, policy=jax.checkpoint_policies.checkpoint_dots)\ndef predict(params, x):\n  for W in params[:-1]:\n    x = apply_layer(W, x)\n  return jnp.dot(params[-1], x)\n```\n\nExample:\n```text\n@jax.checkpoint\ndef f(x):\n  a = some_function(jnp.arange(10_000_000))  # `a` does not depend on `x`\n  return a * x\n```\n\nExample:\n```text\n@partial(jax.checkpoint, concrete=True)  # OLD jax.checkpoint API\ndef foo(x, is_training):\n  if is_training:\n    return g(x)\n  else:\n    return h(x)\n```\n\nExample:\n```text\n@partial(jax.checkpoint, static_argnums=(1,))  # NEW jax.checkpoint API\ndef foo(x, is_training):\n  if is_training:\n    ...\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.879Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":4,"totalLines":42,"estimatedTokens":191}}138{"id":"doc-handling_of_closed_over_constants_jax_documentat-751834d1","source":"documentation","title":"Handling of closed-over constants — JAX documentation","url":"https://docs.jax.dev/en/latest/internals/constants.html","text":"Example:\n```text\nimport numpy as np\nfrom jax import jit\nfrom jax import numpy as jnp\n\na_jax_array = jnp.ones((16,), dtype=np.float32)\n\n@jit\ndef f(x):\n  return x + a_jax_array + np.full((16,), 42.) + jnp.full((16,), 142.)\n```\n\nExample:\n```text\nconst = jnp.array([42.])\nf = jax.jit(lambda: const)\n\nf()\nf()\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.880Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":2,"totalLines":23,"estimatedTokens":81}}139{"id":"doc-autodidax_jax_core_from_scratch_jax_documentatio-b8dc7ea7","source":"documentation","title":"Autodidax: JAX core from scratch — JAX documentation","url":"https://docs.jax.dev/en/latest/autodidax.html","text":"Example:\n```text\ndef f(x):\n  y = sin(x) * 2.\n  z = - y + x\n  return z\n```\n\nExample:\n```text\nfrom typing import NamedTuple\n\nclass Primitive(NamedTuple):\n  name: str\n\nadd_p = Primitive('add')\nmul_p = Primitive('mul')\nneg_p = Primitive(\"neg\")\nsin_p = Primitive(\"sin\")\ncos_p = Primitive(\"cos\")\nreduce_sum_p = Primitive(\"reduce_sum\")\ngreater_p = Primitive(\"greater\")\nless_p = Primitive(\"less\")\ntranspose_p = Primitive(\"transpose\")\nbroadcast_p = Primitive(\"broadcast\")\n\ndef add(x, y): return bind1(add_p, x, y)\ndef mul(x, y): return bind1(mul_p, x, y)\ndef neg(x): return bind1(neg_p, x)\ndef sin(x): return bind1(sin_p, x)\ndef cos(x): return bind1(cos_p, x)\ndef greater(x, y): return bind1(greater_p, x, y)\ndef less(x, y): return bind1(less_p, x, y)\ndef transpose(x, perm): return bind1(transpose_p, x, perm=perm)\ndef broadcast(x, shape, axes): return bind1(broadcast_p, x, shape=shape, axes=axes)\ndef reduce_sum(x, axis=None):\n  if axis is None:\n    axis = tuple(range(np.ndim(x)))\n  if type(axis) is int:\n    axis = (axis,)\n  return bind1(reduce_sum_p, x, axis=axis)\n\ndef bind1(prim, *args, **params):\n  out, = bind(prim, *args, **params)\n  return out\n```\n\nExample:\n```text\nfrom collections.abc import Sequence\nfrom contextlib import contextmanager\nfrom typing import Any\n\nclass MainTrace(NamedTuple):\n  level: int\n  trace_type: type['Trace']\n  global_data: Any | None\n\ntrace_stack: list[MainTrace] = []\ndynamic_trace: MainTrace | None = None  # to be employed in Part 3\n\n@contextmanager\ndef new_main(trace_type: type['Trace'], global_data=None):\n  level = len(trace_stack)\n  main = MainTrace(level, trace_type, global_data)\n  trace_stack.append(main)\n\n  try:\n    yield main\n  finally:\n    trace_stack.pop()\n```\n\nExample:\n```text\nclass Trace:\n  main: MainTrace\n\n  def __init__(self, main: MainTrace) -> None:\n    self.main = main\n\n  def pure(self, val): assert False  # must override\n  def lift(self, val): assert False  # must override\n\n  def process_primitive(self, primitive, tracers, params):\n    assert False  # must override\n```\n\nExample:\n```text\nimport numpy as np\n\nclass Tracer:\n  _trace: Trace\n\n  __array_priority__ = 1000\n\n  @property\n  def aval(self):\n    assert False  # must override\n\n  def full_lower(self):\n    return self  # default implementation\n\n  def __neg__(self): return self.aval._neg(self)\n  def __add__(self, other): return self.aval._add(self, other)\n  def __radd__(self, other): return self.aval._radd(self, other)\n  def __mul__(self, other): return self.aval._mul(self, other)\n  def __rmul__(self, other): return self.aval._rmul(self, other)\n  def __gt__(self, other): return self.aval._gt(self, other)\n  def __lt__(self, other): return self.aval._lt(self, other)\n  def __bool__(self): return self.aval._bool(self)\n  def __nonzero__(self): return self.aval._nonzero(self)\n\n  def __getattr__(self, name):\n    try:\n      return getattr(self.aval, name)\n    except AttributeError:\n      raise AttributeError(f\"{self.__class__.__name__} has no attribute {name}\")\n\ndef swap(f): return lambda x, y: f(y, x)\n```\n\nExample:\n```text\nclass ShapedArray:\n  array_abstraction_level = 1\n  shape: tuple[int, ...]\n  dtype: np.dtype\n  weak_type: bool = False\n\n  def __init__(self, shape, dtype):\n    self.shape = shape\n    self.dtype = dtype\n\n  @property\n  def ndim(self):\n    return len(self.shape)\n\n  _neg = staticmethod(neg)\n  _add = staticmethod(add)\n  _radd = staticmethod(swap(add))\n  _mul = staticmethod(mul)\n  _rmul = staticmethod(swap(mul))\n  _gt = staticmethod(greater)\n  _lt = staticmethod(less)\n\n  @staticmethod\n  def _bool(tracer):\n    raise Exception(\"ShapedArray can't be unambiguously converted to bool\")\n\n  @staticmethod\n  def _nonzero(tracer):\n    raise Exception(\"ShapedArray can't be unambiguously converted to bool\")\n\n  def str_short(self):\n    return f'{self.dtype.name}[{\",\".join(str(d) for d in self.shape)}]'\n\n  def __hash__(self):\n    return hash((self.shape, self.dtype))\n\n  def __eq__(self, other):\n    return (type(self) is type(other) and\n            self.shape == other.shape and self.dtype == other.dtype)\n\n  def __repr__(self):\n    return f\"ShapedArray(shape={self.shape}, dtype={self.dtype})\"\n\nclass ConcreteArray(ShapedArray):\n  array_abstraction_level = 2\n  val: np.ndarray\n\n  def __init__(self, val):\n    self.val = val\n    self.shape = val.shape\n    self.dtype = val.dtype\n\n  @staticmethod\n  def _bool(tracer):\n    return bool(tracer.aval.val)\n\n  @staticmethod\n  def _nonzero(tracer):\n    return bool(tracer.aval.val)\n\ndef get_aval(x):\n  if isinstance(x, Tracer):\n    return x.aval\n  elif type(x) in jax_types:\n    return ConcreteArray(np.asarray(x))\n  else:\n    raise TypeError(x)\n\njax_types = {bool, int, float,\n             np.bool_, np.int32, np.int64, np.float32, np.float64, np.ndarray}\n```\n\nExample:\n```text\ndef bind(prim, *args, **params):\n  top_trace = find_top_trace(args)\n  tracers = [full_raise(top_trace, arg) for arg in args]\n  outs = top_trace.process_primitive(prim, tracers, params)\n  return [full_lower(out) for out in outs]\n```\n\nExample:\n```text\nimport operator as op\n\ndef find_top_trace(xs) -> Trace:\n  top_main = max((x._trace.main for x in xs if isinstance(x, Tracer)),\n                 default=trace_stack[0], key=op.attrgetter('level'))\n  if dynamic_trace and dynamic_trace.level > top_main.level:\n    top_main = dynamic_trace\n  return top_main.trace_type(top_main)\n```\n\nExample:\n```text\ndef full_lower(val: Any):\n  if isinstance(val, Tracer):\n    return val.full_lower()\n  else:\n    return val\n\ndef full_raise(trace: Trace, val: Any) -> Tracer:\n  if not isinstance(val, Tracer):\n    assert type(val) in jax_types\n    return trace.pure(val)\n  level = trace.main.level\n  if val._trace.main is trace.main:\n    return val\n  elif val._trace.main.level < level:\n    return trace.lift(val)\n  elif val._trace.main.level > level:\n    raise Exception(f\"Can't lift level {val._trace.main.level} to {level}.\")\n  else:  # val._trace.level == level\n    raise Exception(f\"Different traces at same level: {val._trace}, {trace}.\")\n```\n\nExample:\n```text\nclass EvalTrace(Trace):\n  pure = lift = lambda self, x: x  # no boxing in Tracers needed\n\n  def process_primitive(self, primitive, tracers, params):\n    return impl_rules[primitive](*tracers, **params)\n\ntrace_stack.append(MainTrace(0, EvalTrace, None))  # special bottom of the stack\n\n# NB: in JAX, instead of a dict we attach impl rules to the Primitive instance\nimpl_rules = {}\n\nimpl_rules[add_p] = lambda x, y: [np.add(x, y)]\nimpl_rules[mul_p] = lambda x, y: [np.multiply(x, y)]\nimpl_rules[neg_p] = lambda x: [np.negative(x)]\nimpl_rules[sin_p] = lambda x: [np.sin(x)]\nimpl_rules[cos_p] = lambda x: [np.cos(x)]\nimpl_rules[reduce_sum_p] = lambda x, *, axis: [np.sum(x, axis)]\nimpl_rules[greater_p] = lambda x, y: [np.greater(x, y)]\nimpl_rules[less_p] = lambda x, y: [np.less(x, y)]\nimpl_rules[transpose_p] = lambda x, *, perm: [np.transpose(x, perm)]\n\ndef broadcast_impl(x, *, shape, axes):\n  for axis in sorted(axes):\n    x = np.expand_dims(x, axis)\n  return [np.broadcast_to(x, shape)]\nimpl_rules[broadcast_p] = broadcast_impl\n```\n\nExample:\n```text\ndef f(x):\n  y = sin(x) * 2.\n  z = - y + x\n  return z\n\nprint(f(3.0))\n```\n\nExample:\n```text\n2.7177599838802657\n```\n\nExample:\n```text\nimport builtins\n\ndef zeros_like(val):\n  aval = get_aval(val)\n  return np.zeros(aval.shape, aval.dtype)\n\ndef unzip2(pairs):\n  lst1, lst2 = [], []\n  for x1, x2 in pairs:\n    lst1.append(x1)\n    lst2.append(x2)\n  return lst1, lst2\n\ndef map(f, *xs):\n  return list(builtins.map(f, *xs))\n\ndef zip(*args):\n  fst, *rest = args = map(list, args)\n  n = len(fst)\n  for arg in rest:\n    assert len(arg) == n\n  return list(builtins.zip(*args))\n```\n\nExample:\n```text\nclass JVPTracer(Tracer):\n  def __init__(self, trace, primal, tangent):\n    self._trace = trace\n    self.primal = primal\n    self.tangent = tangent\n\n  @property\n  def aval(self):\n    return get_aval(self.primal)\n\nclass JVPTrace(Trace):\n  pure = lift = lambda self, val: JVPTracer(self, val, zeros_like(val))\n\n  def process_primitive(self, primitive, tracers, params):\n    primals_in, tangents_in = unzip2((t.primal, t.tangent) for t in tracers)\n    jvp_rule = jvp_rules[primitive]\n    primal_outs, tangent_outs = jvp_rule(primals_in, tangents_in, **params)\n    return [JVPTracer(self, x, t) for x, t in zip(primal_outs, tangent_outs)]\n\njvp_rules = {}\n```\n\nExample:\n```text\ndef add_jvp(primals, tangents):\n  (x, y), (x_dot, y_dot) = primals, tangents\n  return [x + y], [x_dot + y_dot]\njvp_rules[add_p] = add_jvp\n\ndef mul_jvp(primals, tangents):\n  (x, y), (x_dot, y_dot) = primals, tangents\n  return [x * y], [x_dot * y + x * y_dot]\njvp_rules[mul_p] = mul_jvp\n\ndef sin_jvp(primals, tangents):\n  (x,), (x_dot,) = primals, tangents\n  return [sin(x)], [cos(x) * x_dot]\njvp_rules[sin_p] = sin_jvp\n\ndef cos_jvp(primals, tangents):\n  (x,), (x_dot,) = primals, tangents\n  return [cos(x)], [-sin(x) * x_dot]\njvp_rules[cos_p] = cos_jvp\n\ndef neg_jvp(primals, tangents):\n  (x,), (x_dot,) = primals, tangents\n  return [neg(x)], [neg(x_dot)]\njvp_rules[neg_p] = neg_jvp\n\ndef reduce_sum_jvp(primals, tangents, *, axis):\n  (x,), (x_dot,) = primals, tangents\n  return [reduce_sum(x, axis)], [reduce_sum(x_dot, axis)]\njvp_rules[reduce_sum_p] = reduce_sum_jvp\n\ndef greater_jvp(primals, tangents):\n  (x, y), _ = primals, tangents\n  out_primal = greater(x, y)\n  return [out_primal], [zeros_like(out_primal)]\njvp_rules[greater_p] = greater_jvp\n\ndef less_jvp(primals, tangents):\n  (x, y), _ = primals, tangents\n  out_primal = less(x, y)\n  return [out_primal], [zeros_like(out_primal)]\njvp_rules[less_p] = less_jvp\n```\n\nExample:\n```text\ndef jvp_v1(f, primals, tangents):\n  with new_main(JVPTrace) as main:\n    trace = JVPTrace(main)\n    tracers_in = [JVPTracer(trace, x, t) for x, t in zip(primals, tangents)]\n    out = f(*tracers_in)\n    tracer_out = full_raise(trace, out)\n    primal_out, tangent_out = tracer_out.primal, tracer_out.tangent\n  return primal_out, tangent_out\n```\n\nExample:\n```text\nx = 3.0\ny, sin_deriv_at_3 = jvp_v1(sin, (x,), (1.0,))\nprint(sin_deriv_at_3)\nprint(cos(3.0))\n```\n\nExample:\n```text\n-0.9899924966004454\n-0.9899924966004454\n```\n\nExample:\n```text\ndef f(x):\n  y = sin(x) * 2.\n  z = - y + x\n  return z\n\nx, xdot = 3., 1.\ny, ydot = jvp_v1(f, (x,), (xdot,))\nprint(y)\nprint(ydot)\n```\n\nExample:\n```text\n2.7177599838802657\n2.979984993200891\n```\n\nExample:\n```text\ndef deriv(f):\n  return lambda x: jvp_v1(f, (x,), (1.,))[1]\n\nprint(deriv(sin)(3.))\nprint(deriv(deriv(sin))(3.))\nprint(deriv(deriv(deriv(sin)))(3.))\nprint(deriv(deriv(deriv(deriv(sin))))(3.))\n```\n\nExample:\n```text\n-0.9899924966004454\n-0.1411200080598672\n0.9899924966004454\n0.1411200080598672\n```\n\nExample:\n```text\ndef f(x):\n  if x > 0.:  # Python control flow\n    return 2. * x\n  else:\n    return x\n\nprint(deriv(f)(3.))\nprint(deriv(f)(-3.))\n```\n\nExample:\n```text\n2.0\n1.0\n```\n\nExample:\n```text\ndef jvp_flat(f, primals, tangents):\n  with new_main(JVPTrace) as main:\n    trace = JVPTrace(main)\n    tracers_in = [JVPTracer(trace, x, t) for x, t in zip(primals, tangents)]\n    outs = f(*tracers_in)\n    tracers_out = [full_raise(trace, out) for out in outs]\n    primals_out, tangents_out = unzip2((t.primal, t.tangent) for t in tracers_out)\n  return primals_out, tangents_out\n```\n\nExample:\n```text\ndef jvp(f, primals, tangents):\n  primals_flat, in_tree = tree_flatten(primals)\n  tangents_flat, in_tree2 = tree_flatten(tangents)\n  if in_tree != in_tree2: raise TypeError\n  f, out_tree = flatten_fun(f, in_tree)\n  primals_out_flat, tangents_out_flat = jvp_flat(f, primals_flat, tangents_flat)\n  primals_out = tree_unflatten(out_tree(), primals_out_flat)\n  tangents_out = tree_unflatten(out_tree(), tangents_out_flat)\n  return primals_out, tangents_out\n```\n\nExample:\n```text\ndef flatten_fun(f, in_tree):\n  store = Store()\n\n  def flat_fun(*args_flat):\n    pytree_args = tree_unflatten(in_tree, args_flat)\n    out = f(*pytree_args)\n    out_flat, out_tree = tree_flatten(out)\n    store.set_value(out_tree)\n    return out_flat\n\n  return flat_fun, store\n\nclass Empty: pass\nempty = Empty()\n\nclass Store:\n  val = empty\n\n  def set_value(self, val):\n    assert self.val is empty\n    self.val = val\n\n  def __call__(self):\n    return self.val\n```\n\nExample:\n```text\nfrom collections.abc import Hashable, Iterable, Iterator\nimport itertools as it\nfrom collections.abc import Callable\n\nclass NodeType(NamedTuple):\n  name: str\n  to_iterable: Callable\n  from_iterable: Callable\n\ndef register_pytree_node(ty: type, to_iter: Callable, from_iter: Callable\n                         ) -> None:\n  node_types[ty] = NodeType(str(ty), to_iter, from_iter)\n\nnode_types: dict[type, NodeType] = {}\nregister_pytree_node(tuple, lambda t: (None, t), lambda _, xs: tuple(xs))\nregister_pytree_node(list,  lambda l: (None, l), lambda _, xs:  list(xs))\nregister_pytree_node(dict,\n                     lambda d: map(tuple, unzip2(sorted(d.items()))),\n                     lambda keys, vals: dict(zip(keys, vals)))\n\nclass PyTreeDef(NamedTuple):\n  node_type: NodeType\n  node_metadata: Hashable\n  child_treedefs: tuple['PyTreeDef', ...]\n\nclass Leaf: pass\nleaf = Leaf()\n\ndef tree_flatten(x: Any) -> tuple[list[Any], PyTreeDef]:\n  children_iter, treedef = _tree_flatten(x)\n  return list(children_iter), treedef\n\ndef _tree_flatten(x: Any) -> tuple[Iterable, PyTreeDef]:\n  node_type = node_types.get(type(x))\n  if node_type:\n    node_metadata, children = node_type.to_iterable(x)\n    children_flat, child_trees = unzip2(map(_tree_flatten, children))\n    flattened = it.chain.from_iterable(children_flat)\n    return flattened, PyTreeDef(node_type, node_metadata, tuple(child_trees))\n  else:\n    return [x], leaf\n\ndef tree_unflatten(treedef: PyTreeDef, xs: list[Any]) -> Any:\n  return _tree_unflatten(treedef, iter(xs))\n\ndef _tree_unflatten(treedef: PyTreeDef, xs: Iterator) -> Any:\n  if treedef is leaf:\n    return next(xs)\n  else:\n    children = (_tree_unflatten(t, xs) for t in treedef.child_treedefs)\n    return treedef.node_type.from_iterable(treedef.node_metadata, children)\n```\n\nExample:\n```text\ndef f(x):\n  y = sin(x) * 2.\n  z = - y + x\n  return {'hi': z, 'there': [x, y]}\n\nx, xdot = 3., 1.\ny, ydot = jvp(f, (x,), (xdot,))\nprint(y)\nprint(ydot)\n```\n\nExample:\n```text\n{'hi': np.float64(2.7177599838802657), 'there': [3.0, np.float64(0.2822400161197344)]}\n{'hi': np.float64(2.979984993200891), 'there': [1.0, np.float64(-1.9799849932008908)]}\n```\n\nExample:\n```text\ndef mapped_aval(batch_dim, aval):\n  shape = list(aval.shape)\n  del shape[batch_dim]\n  return ShapedArray(tuple(shape), aval.dtype)\n\ndef move_batch_axis(axis_size, src, dst, x):\n  if src is not_mapped:\n    target_shape = list(np.shape(x))\n    target_shape.insert(dst, axis_size)\n    return broadcast(x, target_shape, [dst])\n  elif src == dst:\n    return x\n  else:\n    return moveaxis(x, src, dst)\n\ndef moveaxis(x, src: int, dst: int):\n  perm = [i for i in range(np.ndim(x)) if i != src]\n  perm.insert(dst, src)\n  return transpose(x, perm)\n```\n\nExample:\n```text\nfrom typing import Union\n\nclass NotMapped: pass\nnot_mapped = NotMapped()\n\nBatchAxis = Union[NotMapped, int]\n\nclass BatchTracer(Tracer):\n  def __init__(self, trace, val, batch_dim: BatchAxis):\n    self._trace = trace\n    self.val = val\n    self.batch_dim = batch_dim\n\n  @property\n  def aval(self):\n    if self.batch_dim is not_mapped:\n      return get_aval(self.val)\n    else:\n      return mapped_aval(self.batch_dim, get_aval(self.val))\n\n  def full_lower(self):\n    if self.batch_dim is not_mapped:\n      return full_lower(self.val)\n    else:\n      return self\n\nclass BatchTrace(Trace):\n  pure = lift = lambda self, val: BatchTracer(self, val, not_mapped)\n\n  def process_primitive(self, primitive, tracers, params):\n    vals_in, bdims_in = unzip2((t.val, t.batch_dim) for t in tracers)\n    vmap_rule = vmap_rules[primitive]\n    val_outs, bdim_outs = vmap_rule(self.axis_size, vals_in, bdims_in, **params)\n    return [BatchTracer(self, x, bd) for x, bd in zip(val_outs, bdim_outs)]\n\n  @property\n  def axis_size(self):\n    return self.main.global_data\n\nvmap_rules = {}\n```\n\nExample:\n```text\nfrom functools import partial\n\ndef binop_batching_rule(op, axis_size, vals_in, dims_in):\n  (x, y), (x_bdim, y_bdim) = vals_in, dims_in\n  if x_bdim != y_bdim:\n    if x_bdim is not_mapped:\n      x = move_batch_axis(axis_size, x_bdim, y_bdim, x)\n      x_bdim = y_bdim\n    else:\n      y = move_batch_axis(axis_size, y_bdim, x_bdim, y)\n  return [op(x, y)], [x_bdim]\nvmap_rules[add_p] = partial(binop_batching_rule, add)\nvmap_rules[mul_p] = partial(binop_batching_rule, mul)\n\ndef vectorized_unop_batching_rule(op, axis_size, vals_in, dims_in):\n  (x,), (x_bdim,) = vals_in, dims_in\n  return [op(x)], [x_bdim]\nvmap_rules[sin_p] = partial(vectorized_unop_batching_rule, sin)\nvmap_rules[cos_p] = partial(vectorized_unop_batching_rule, cos)\nvmap_rules[neg_p] = partial(vectorized_unop_batching_rule, neg)\n\ndef reduce_sum_batching_rule(axis_size, vals_in, dims_in, *, axis):\n  (x,), (x_bdim,) = vals_in, dims_in\n  new_axis = tuple(ax + (x_bdim <= ax) for ax in axis)\n  out_bdim = x_bdim - sum(ax < x_bdim for ax in axis)\n  return [reduce_sum(x, new_axis)], [out_bdim]\nvmap_rules[reduce_sum_p] = reduce_sum_batching_rule\n```\n\nExample:\n```text\ndef vmap_flat(f, in_axes, *args):\n  axis_size, = {x.shape[ax] for x, ax in zip(args, in_axes)\n                if ax is not not_mapped}\n  with new_main(BatchTrace, axis_size) as main:\n    trace = BatchTrace(main)\n    tracers_in = [BatchTracer(trace, x, ax) if ax is not None else x\n                  for x, ax in zip(args, in_axes)]\n    outs = f(*tracers_in)\n    tracers_out = [full_raise(trace, out) for out in outs]\n    vals_out, bdims_out = unzip2((t.val, t.batch_dim) for t in tracers_out)\n  outs_transposed = [move_batch_axis(axis_size, bdim, 0, val_out)\n                     for val_out, bdim in zip(vals_out, bdims_out)]\n  return outs_transposed\n\ndef vmap(f, in_axes):\n  def batched_f(*args):\n    args_flat, in_tree = tree_flatten(args)\n    in_axes_flat, in_tree2 = tree_flatten(in_axes)\n    if in_tree != in_tree2: raise TypeError\n    f_flat, out_tree = flatten_fun(f, in_tree)\n    outs_flat = vmap_flat(f_flat, in_axes_flat, *args_flat)\n    return tree_unflatten(out_tree(), outs_flat)\n  return batched_f\n```\n\nExample:\n```text\ndef add_one_to_a_scalar(scalar):\n  assert np.ndim(scalar) == 0\n  return 1 + scalar\n\nvector_in = np.arange(3.)\nvector_out = vmap(add_one_to_a_scalar, (0,))(vector_in)\n\nprint(vector_in)\nprint(vector_out)\n```\n\nExample:\n```text\n[0. 1. 2.]\n[1. 2. 3.]\n```\n\nExample:\n```text\ndef jacfwd(f, x):\n  pushfwd = lambda v: jvp(f, (x,), (v,))[1]\n  vecs_in = np.eye(np.size(x)).reshape(np.shape(x) * 2)\n  return vmap(pushfwd, (0,))(vecs_in)\n\ndef f(x):\n  return sin(x)\n\njacfwd(f, np.arange(3.))\n```\n\nExample:\n```text\narray([[ 1.        ,  0.        , -0.        ],\n       [ 0.        ,  0.54030231, -0.        ],\n       [ 0.        ,  0.        , -0.41614684]])\n```\n\nExample:\n```text\njaxpr ::=\n  { lambda <binder> , ... .\n    let <eqn>\n        ...\n    in ( <atom> , ... ) }\n\nbinder ::= <var>:<array_type>\nvar ::= a | b | c | ...\natom ::= <var> | <literal>\nliteral ::= <int32> | <int64> | <float32> | <float64>\n\neqn ::= <binder> , ... = <primitive> [ <params> ] <atom> , ...\n```\n\nExample:\n```text\njaxpr_type ::= [ <array_type> , ... ] -> [ <array_type> , ... ]\narray_type ::= <dtype>[<shape>]\ndtype ::= f32 | f64 | i32 | i64\nshape ::= <int> , ...\n```\n\nExample:\n```text\nclass Var:\n  aval: ShapedArray\n  def __init__(self, aval): self.aval = aval\n\nclass Lit:\n  val: Any\n  aval: ShapedArray\n\n  def __init__(self, val):\n    self.aval = aval = raise_to_shaped(get_aval(val))\n    self.val = np.array(val, aval.dtype)\n\nAtom = Union[Var, Lit]\n\nclass JaxprEqn(NamedTuple):\n  primitive: Primitive\n  inputs: list[Atom]\n  params: dict[str, Any]\n  out_binders: list[Var]\n\nclass Jaxpr(NamedTuple):\n  in_binders: list[Var]\n  eqns: list[JaxprEqn]\n  outs: list[Atom]\n\n  def __hash__(self): return id(self)\n  __eq__ = op.is_\n\ndef raise_to_shaped(aval):\n  return ShapedArray(aval.shape, aval.dtype)\n```\n\nExample:\n```text\nclass JaxprType(NamedTuple):\n  in_types:  list[ShapedArray]\n  out_types: list[ShapedArray]\n\n  def __repr__(self):\n    in_types = ', '.join(aval.str_short() for aval in self.in_types)\n    out_types = ', '.join(aval.str_short() for aval in self.out_types)\n    return f'({in_types}) -> ({out_types})'\n\ndef typecheck_jaxpr(jaxpr: Jaxpr) -> JaxprType:\n  env: set[Var] = set()\n\n  for v in jaxpr.in_binders:\n    if v in env: raise TypeError\n    env.add(v)\n\n  for eqn in jaxpr.eqns:\n    in_types = [typecheck_atom(env, x) for x in eqn.inputs]\n    out_types = abstract_eval_rules[eqn.primitive](*in_types, **eqn.params)\n    for out_binder, out_type in zip(eqn.out_binders, out_types):\n      if not out_type == out_binder.aval: raise TypeError\n    for out_binder in eqn.out_binders:\n      if out_binder in env: raise TypeError\n      env.add(out_binder)\n\n  in_types = [v.aval for v in jaxpr.in_binders]\n  out_types = [typecheck_atom(env, x) for x in jaxpr.outs]\n  return JaxprType(in_types, out_types)\n\ndef typecheck_atom(env: set[Var], x: Atom) -> ShapedArray:\n  if isinstance(x, Var):\n    if x not in env: raise TypeError(\"unbound variable\")\n    return x.aval\n  elif isinstance(x, Lit):\n    return raise_to_shaped(get_aval(x.val))\n  else:\n    assert False\n```\n\nExample:\n```text\ndef eval_jaxpr(jaxpr: Jaxpr, args: list[Any]) -> list[Any]:\n  env: dict[Var, Any] = {}\n\n  def read(x: Atom) -> Any:\n    return env[x] if type(x) is Var else x.val\n\n  def write(v: Var, val: Any) -> None:\n    assert v not in env  # single-assignment\n    env[v] = val\n\n  map(write, jaxpr.in_binders, args)\n  for eqn in jaxpr.eqns:\n    in_vals = map(read, eqn.inputs)\n    outs = bind(eqn.primitive, *in_vals, **eqn.params)\n    map(write, eqn.out_binders, outs)\n  return map(read, jaxpr.outs)\n\ndef jaxpr_as_fun(jaxpr: Jaxpr):\n  return lambda *args: eval_jaxpr(jaxpr, args)\n```\n\nExample:\n```text\ndef split_list(lst: list[Any], n: int) -> tuple[list[Any], list[Any]]:\n  assert 0 <= n <= len(lst)\n  return lst[:n], lst[n:]\n\ndef partition_list(bs: list[bool], l: list[Any]) -> tuple[list[Any], list[Any]]:\n  assert len(bs) == len(l)\n  lists = lst1, lst2 = [], []\n  for b, x in zip(bs, l):\n    lists[b].append(x)\n  return lst1, lst2\n```\n\nExample:\n```text\n# NB: the analogous class in JAX is called 'DynamicJaxprTracer'\nclass JaxprTracer(Tracer):\n  __slots__ = ['aval']\n  aval: ShapedArray\n\n  def __init__(self, trace, aval):\n    self._trace = trace\n    self.aval = aval\n\n# NB: the analogous class in JAX is called 'DynamicJaxprTrace'\nclass JaxprTrace(Trace):\n  def new_arg(self, aval: ShapedArray) -> JaxprTracer:\n    aval = raise_to_shaped(aval)\n    tracer = self.builder.new_tracer(self, aval)\n    self.builder.tracer_to_var[id(tracer)] = Var(aval)\n    return tracer\n\n  def get_or_make_const_tracer(self, val: Any) -> JaxprTracer:\n    tracer = self.builder.const_tracers.get(id(val))\n    if tracer is None:\n      tracer = self.builder.new_tracer(self, raise_to_shaped(get_aval(val)))\n      self.builder.add_const(tracer, val)\n    return tracer\n  pure = lift = get_or_make_const_tracer\n\n  def process_primitive(self, primitive, tracers, params):\n    avals_in = [t.aval for t in tracers]\n    avals_out = abstract_eval_rules[primitive](*avals_in, **params)\n    out_tracers = [self.builder.new_tracer(self, a) for a in avals_out]\n    inputs = [self.builder.getvar(t) for t in tracers]\n    outvars = [self.builder.add_var(t) for t in out_tracers]\n    self.builder.add_eqn(JaxprEqn(primitive, inputs, params, outvars))\n    return out_tracers\n\n  @property\n  def builder(self):\n    return self.main.global_data\n\n# NB: in JAX, we instead attach abstract eval rules to Primitive instances\nabstract_eval_rules = {}\n```\n\nExample:\n```text\nclass JaxprBuilder:\n  eqns: list[JaxprEqn]\n  tracer_to_var: dict[int, Var]\n  const_tracers: dict[int, JaxprTracer]\n  constvals: dict[Var, Any]\n  tracers: list[JaxprTracer]\n\n  def __init__(self):\n    self.eqns = []\n    self.tracer_to_var = {}\n    self.const_tracers = {}\n    self.constvals = {}\n    self.tracers = []\n\n  def new_tracer(self, trace: JaxprTrace, aval: ShapedArray) -> JaxprTracer:\n    tracer = JaxprTracer(trace, aval)\n    self.tracers.append(tracer)\n    return tracer\n\n  def add_eqn(self, eqn: JaxprEqn) -> None:\n    self.eqns.append(eqn)\n\n  def add_var(self, tracer: JaxprTracer) -> Var:\n    assert id(tracer) not in self.tracer_to_var\n    var = self.tracer_to_var[id(tracer)] = Var(tracer.aval)\n    return var\n\n  def getvar(self, tracer: JaxprTracer) -> Var:\n    var = self.tracer_to_var.get(id(tracer))\n    assert var is not None\n    return var\n\n  def add_const(self, tracer: JaxprTracer, val: Any) -> Var:\n    var = self.add_var(tracer)\n    self.const_tracers[id(val)] = tracer\n    self.constvals[var] = val\n    return var\n\n  def build(self, in_tracers: list[JaxprTracer], out_tracers: list[JaxprTracer]\n            ) -> tuple[Jaxpr, list[Any]]:\n    constvars, constvals = unzip2(self.constvals.items())\n    t2v = lambda t: self.tracer_to_var[id(t)]\n    in_binders = constvars + [t2v(t) for t in in_tracers]\n    out_vars = [t2v(t) for t in out_tracers]\n    jaxpr = Jaxpr(in_binders, self.eqns, out_vars)\n    typecheck_jaxpr(jaxpr)\n    jaxpr, constvals = _inline_literals(jaxpr, constvals)\n    return jaxpr, constvals\n```\n\nExample:\n```text\ndef _inline_literals(jaxpr: Jaxpr, consts: list[Any]) -> tuple[Jaxpr, list[Any]]:\n  const_binders, other_binders = split_list(jaxpr.in_binders, len(consts))\n  scalars = [type(x) in jax_types and not get_aval(x).shape for x in consts]\n  new_const_binders, lit_binders = partition_list(scalars, const_binders)\n  new_consts, lit_vals = partition_list(scalars, consts)\n  literals = dict(zip(lit_binders, map(Lit, lit_vals)))\n  new_eqns = [JaxprEqn(eqn.primitive, [literals.get(x, x) for x in eqn.inputs],\n                       eqn.params, eqn.out_binders) for eqn in jaxpr.eqns]\n  new_outs = [literals.get(x, x) for x in jaxpr.outs]\n  new_jaxpr = Jaxpr(new_const_binders + other_binders, new_eqns, new_outs)\n  typecheck_jaxpr(new_jaxpr)\n  return new_jaxpr, new_consts\n```\n\nExample:\n```text\ndef binop_abstract_eval(x: ShapedArray, y: ShapedArray) -> list[ShapedArray]:\n  if not isinstance(x, ShapedArray) or not isinstance(y, ShapedArray):\n    raise TypeError\n  if raise_to_shaped(x) != raise_to_shaped(y): raise TypeError\n  return [ShapedArray(x.shape, x.dtype)]\n\nabstract_eval_rules[add_p] = binop_abstract_eval\nabstract_eval_rules[mul_p] = binop_abstract_eval\n\ndef compare_abstract_eval(x: ShapedArray, y: ShapedArray) -> list[ShapedArray]:\n  if not isinstance(x, ShapedArray) or not isinstance(y, ShapedArray):\n    raise TypeError\n  if x.shape != y.shape: raise TypeError\n  return [ShapedArray(x.shape, np.dtype('bool'))]\nabstract_eval_rules[greater_p] = compare_abstract_eval\nabstract_eval_rules[less_p] = compare_abstract_eval\n\ndef vectorized_unop_abstract_eval(x: ShapedArray) -> list[ShapedArray]:\n  return [ShapedArray(x.shape, x.dtype)]\n\nabstract_eval_rules[sin_p] = vectorized_unop_abstract_eval\nabstract_eval_rules[cos_p] = vectorized_unop_abstract_eval\nabstract_eval_rules[neg_p] = vectorized_unop_abstract_eval\n\ndef reduce_sum_abstract_eval(x: ShapedArray, *, axis: tuple[int, ...]\n                             ) -> list[ShapedArray]:\n  axis_ = set(axis)\n  new_shape = [d for i, d in enumerate(x.shape) if i not in axis_]\n  return [ShapedArray(tuple(new_shape), x.dtype)]\nabstract_eval_rules[reduce_sum_p] = reduce_sum_abstract_eval\n\ndef broadcast_abstract_eval(x: ShapedArray, *, shape: Sequence[int],\n                            axes: Sequence[int]) -> list[ShapedArray]:\n  return [ShapedArray(tuple(shape), x.dtype)]\nabstract_eval_rules[broadcast_p] = broadcast_abstract_eval\n```\n\nExample:\n```text\nfrom functools import lru_cache\n\n@lru_cache  # ShapedArrays are hashable\ndef make_jaxpr_v1(f, *avals_in):\n  avals_in, in_tree = tree_flatten(avals_in)\n  f, out_tree = flatten_fun(f, in_tree)\n\n  builder = JaxprBuilder()\n  with new_main(JaxprTrace, builder) as main:\n    trace = JaxprTrace(main)\n    tracers_in = [trace.new_arg(aval) for aval in avals_in]\n    outs = f(*tracers_in)\n    tracers_out = [full_raise(trace, out) for out in outs]\n    jaxpr, consts = builder.build(tracers_in, tracers_out)\n  return jaxpr, consts, out_tree()\n```\n\nExample:\n```text\nfrom collections import defaultdict\nimport string\n\nclass PPrint:\n  lines: list[tuple[int, str]]\n\n  def __init__(self, lines):\n    self.lines = lines\n\n  def indent(self, indent: int) -> 'PPrint':\n    return PPrint([(indent + orig_indent, s) for orig_indent, s in self.lines])\n\n  def __add__(self, rhs: 'PPrint') -> 'PPrint':\n    return PPrint(self.lines + rhs.lines)\n\n  def __rshift__(self, rhs: 'PPrint') -> 'PPrint':\n    if not rhs.lines: return self\n    if not self.lines: return rhs\n    indent, s = self.lines[-1]\n    indented_block = rhs.indent(indent + len(s))\n    common_line = s + ' ' * rhs.lines[0][0] + rhs.lines[0][1]\n    return PPrint(self.lines[:-1]\n                  + [(indent, common_line)]\n                  + indented_block.lines[1:])\n\n  def __str__(self) -> str:\n    return '\\n'.join(' ' * indent + s for indent, s in self.lines)\n\ndef pp(s: Any) -> PPrint:\n  return PPrint([(0, line) for line in str(s).splitlines()])\n\ndef vcat(ps: list[PPrint]) -> PPrint:\n  return sum(ps, pp(''))\n\ndef pp_jaxpr(jaxpr: Jaxpr) -> PPrint:\n  namegen = (''.join(s) for r in it.count(1)\n             for s in it.permutations(string.ascii_lowercase, r))\n  names = defaultdict(lambda: next(namegen))\n  in_binders = ', '.join(var_str(names, x) for x in jaxpr.in_binders)\n  eqns = vcat([pp_eqn(names, e) for e in jaxpr.eqns])\n  outs = ', '.join(names[v] if isinstance(v, Var) else str(v.val)\n                   for v in jaxpr.outs)\n  return (pp(f'{{ lambda {in_binders} .') +\n          ((pp('let ') >> eqns) + pp(f'in ( {outs} ) }}')).indent(2))\n\ndef var_str(names: defaultdict[Var, str], v: Var) -> str:\n  return f'{names[v]}:{v.aval.str_short()}'\n\ndef pp_eqn(names: defaultdict[Var, str], eqn: JaxprEqn) -> PPrint:\n  rule = pp_rules.get(eqn.primitive)\n  if rule:\n    return rule(names, eqn)\n  else:\n    lhs = pp(' '.join(var_str(names, v) for v in eqn.out_binders))\n    rhs = (pp(eqn.primitive.name) >> pp_params(eqn.params) >>\n           pp(' '.join(names[x] if isinstance(x, Var) else str(x.val)\n                       for x in eqn.inputs)))\n    return lhs >> pp(' = ') >> rhs\n\ndef pp_params(params: dict[str, Any]) -> PPrint:\n  items = sorted(params.items())\n  if items:\n    return pp(' [ ') >> vcat([pp(f'{k}={v}') for k, v in items]) >> pp(' ] ')\n  else:\n    return pp(' ')\n\nJaxpr.__repr__ = lambda self: str(pp_jaxpr(self))\npp_rules: dict[Primitive, Callable[..., PPrint]] = {}\n```\n\nExample:\n```text\njaxpr, consts, _ = make_jaxpr_v1(lambda x: 2. * x, raise_to_shaped(get_aval(3.)))\nprint(jaxpr)\nprint(typecheck_jaxpr(jaxpr))\n```\n\nExample:\n```text\n{ lambda a:float64[] .\n  let b:float64[] = mul 2.0 a\n  in ( b ) }\n(float64[]) -> (float64[])\n```\n\nExample:\n```text\njaxpr, consts, _ = make_jaxpr_v1(lambda: mul(2., 2.))\nprint(jaxpr)\n```\n\nExample:\n```text\n{ lambda  .\n  let \n  in ( 4.0 ) }\n```\n\nExample:\n```text\n@contextmanager\ndef new_dynamic(main: MainTrace):\n  global dynamic_trace\n  prev_dynamic_trace, dynamic_trace = dynamic_trace, main\n  try:\n    yield\n  finally:\n    dynamic_trace = prev_dynamic_trace\n\n@lru_cache\ndef make_jaxpr(f: Callable, *avals_in: ShapedArray,\n               ) -> tuple[Jaxpr, list[Any], PyTreeDef]:\n  avals_in, in_tree = tree_flatten(avals_in)\n  f, out_tree = flatten_fun(f, in_tree)\n\n  builder = JaxprBuilder()\n  with new_main(JaxprTrace, builder) as main:\n    with new_dynamic(main):\n      trace = JaxprTrace(main)\n      tracers_in = [trace.new_arg(aval) for aval in avals_in]\n      outs = f(*tracers_in)\n      tracers_out = [full_raise(trace, out) for out in outs]\n      jaxpr, consts = builder.build(tracers_in, tracers_out)\n  return jaxpr, consts, out_tree()\n\njaxpr, consts, _ = make_jaxpr(lambda: mul(2., 2.))\nprint(jaxpr)\n```\n\nExample:\n```text\n{ lambda  .\n  let a:float64[] = mul 2.0 2.0\n  in ( a ) }\n```\n\nExample:\n```text\ndef jit(f):\n  def f_jitted(*args):\n    avals_in = [raise_to_shaped(get_aval(x)) for x in args]\n    jaxpr, consts, out_tree = make_jaxpr(f, *avals_in)\n    outs = bind(xla_call_p, *consts, *args, jaxpr=jaxpr, num_consts=len(consts))\n    return tree_unflatten(out_tree, outs)\n  return f_jitted\n\nxla_call_p = Primitive('xla_call')\n```\n\nExample:\n```text\nclass IDHashable:\n  val: Any\n\n  def __init__(self, val):\n    self.val = val\n\n  def __hash__(self) -> int:\n    return id(self.val)\n\n  def __eq__(self, other):\n    return type(other) is IDHashable and id(self.val) == id(other.val)\n```\n\nExample:\n```text\nimport io\nfrom jax.extend.mlir import ir\nfrom jax.extend.mlir.dialects import func\nfrom jax.extend.mlir.dialects import stablehlo as hlo\nfrom jax._src import xla_bridge as xb\nfrom jax._src.lib import jax_mlir_ext\nfrom jax._src.lib import xla_client as xc\n\n_dialects_registry = ir.DialectRegistry()\njax_mlir_ext.register_dialects(_dialects_registry, register_pipelines=False)\n\nclass MlirContext(NamedTuple):\n  module: ir.Module\n  symbol_table: ir.SymbolTable\n\ndef xla_call_impl(*args, jaxpr: Jaxpr, num_consts: int):\n  consts, args = args[:num_consts], args[num_consts:]\n  hashable_consts = tuple(map(IDHashable, consts))\n  execute = xla_callable(IDHashable(jaxpr), hashable_consts)\n  return execute(*args)\nimpl_rules[xla_call_p] = xla_call_impl\n\n@lru_cache\ndef xla_callable(hashable_jaxpr: IDHashable,\n                 hashable_consts: tuple[IDHashable, ...]):\n  jaxpr: Jaxpr = hashable_jaxpr.val\n  typecheck_jaxpr(jaxpr)\n  consts = [x.val for x in hashable_consts]\n  in_avals = [v.aval for v in jaxpr.in_binders[len(consts):]]\n\n  with ir.Context() as ctx, ir.Location.unknown(ctx):\n    ctx.append_dialect_registry(_dialects_registry)\n    ctx.load_all_available_dialects()\n    hlo.register_dialect(ctx)\n    m = ir.Module.create()\n    c = MlirContext(m, ir.SymbolTable(m.operation))\n\n    with ir.InsertionPoint(c.module.body):\n      @func.func(*(aval_to_ir_type(aval) for aval in in_avals))\n      def main(*params):\n        return jaxpr_subcomp(c, jaxpr, _hlo_consts(consts) + params)\n\n  output = io.StringIO()\n  c.module.operation.print(file=output)\n  backend = xb.get_backend(None)\n  compiled = backend.compile_and_load(output.getvalue(), backend.devices()[:1])\n  return partial(execute_compiled, compiled, [v.aval for v in jaxpr.outs])\n\ndef _mlir_dtype(dtype: np.dtype) -> ir.Type:\n  if np.issubdtype(dtype, np.signedinteger):\n    return ir.IntegerType.get_signless(np.iinfo(dtype).bits)\n  elif dtype == np.float32:\n    return ir.F32Type.get()\n  elif dtype == np.float64:\n    return ir.F64Type.get()\n  else:\n    raise NotImplementedError(\"MLIR conversion not implemented for \", dtype)\n\ndef aval_to_ir_type(aval: ShapedArray) -> ir.Type:\n  return ir.RankedTensorType.get(aval.shape, _mlir_dtype(aval.dtype))\n\ndef _hlo_const(x: Any) -> ir.Value:\n  a = np.asarray(x)\n  if a.dtype == np.bool_:\n    return hlo.constant(ir.DenseElementsAttr.get(\n      np.array(a, np.bool_), type=ir.IntegerType.get_signless(1),\n      shape=a.shape))\n  else:\n    return hlo.constant(ir.DenseElementsAttr.get(a))\n\ndef _hlo_consts(consts: list[Any]) -> list[ir.Value]:\n  unique_consts = {id(cnst): cnst for cnst in consts}\n  ir_consts = {id_: _hlo_const(cnst) for id_, cnst in unique_consts.items()}\n  return tuple(ir_consts[id(cnst)] for cnst in consts)\n```\n\nExample:\n```text\ndef jaxpr_subcomp(c: MlirContext, jaxpr: Jaxpr, args: list[ir.Value]) -> list[ir.Value]:\n  env: dict[Var, ir.Value] = {}\n\n  def read(x: Atom) -> ir.Value:\n    return env[x] if type(x) is Var else _hlo_const(np.asarray(x.val))\n\n  def write(v: Var, val: ir.Value) -> None:\n    env[v] = val\n\n  map(write, jaxpr.in_binders, args)\n  for eqn in jaxpr.eqns:\n    in_avals = [x.aval for x in eqn.inputs]\n    in_vals = map(read, eqn.inputs)\n    out_avals = [x.aval for x in eqn.out_binders]\n    rule = hlo_translations[eqn.primitive]\n    assert all(isinstance(v, ir.Value) for v in in_vals), in_vals\n    out_vals = rule(c, in_avals, out_avals, in_vals, **eqn.params)\n    assert all(isinstance(v, ir.Value) for v in out_vals), out_vals\n    map(write, eqn.out_binders, out_vals), out_vals\n  return map(read, jaxpr.outs)\n\ndef execute_compiled(compiled, out_avals, *args):\n  input_bufs = [input_handlers[type(x)](x) for x in args]\n  out_bufs = compiled.execute(input_bufs)\n  return [handle_result(aval, buf) for aval, buf in zip(out_avals, out_bufs)]\n\ndef default_input_handler(x):\n  device = xb.get_backend(None).devices()[0]\n  return xc.batched_device_put(get_aval(x),\n                               xc.GSPMDSharding((device,), xc.OpSharding()),\n                               [x], [device], True, enable_x64=True)\n\ninput_handlers = {ty: default_input_handler for ty in\n                  [bool, int, float, np.ndarray, np.float64, np.float32]}\n\ndef handle_result(aval: ShapedArray, buf):\n  del aval  # Unused for now\n  return np.asarray(buf)\n\nhlo_translations = {}\n```\n\nExample:\n```text\ndef direct_translation(op, c, in_avals, out_avals, in_vals):\n  del c, in_avals, out_avals\n  return [op(*in_vals)]\n\nhlo_translations[add_p] = partial(direct_translation, hlo.add)\nhlo_translations[mul_p] = partial(direct_translation, hlo.multiply)\nhlo_translations[neg_p] = partial(direct_translation, hlo.negate)\nhlo_translations[sin_p] = partial(direct_translation, hlo.sine)\nhlo_translations[cos_p] = partial(direct_translation, hlo.cosine)\n\ndef compare_translation(op, c, in_avals, out_avals, in_vals):\n  del c, out_avals\n  return [hlo.compare(*in_vals, hlo.ComparisonDirectionAttr.get(op))]\n\nhlo_translations[greater_p] = partial(compare_translation, \"GT\")\nhlo_translations[less_p] = partial(compare_translation, \"LT\")\n\ndef reduce_sum_translation(c, in_avals, out_avals, in_vals, *, axis):\n  del c\n  (x_aval,), (out_aval,), (x,) = in_avals, out_avals, in_vals\n  op = hlo.ReduceOp(\n    [aval_to_ir_type(out_aval)], [x], [_hlo_const(np.array(0, x_aval.dtype))],\n    axis)\n  scalar_type = aval_to_ir_type(ShapedArray((), x_aval.dtype))\n  reducer_region = op.body.blocks.append(scalar_type, scalar_type)\n  with ir.InsertionPoint(reducer_region):\n    hlo.return_([hlo.add(*reducer_region.arguments)])\n  return op.results\n\nhlo_translations[reduce_sum_p] = reduce_sum_translation\n\ndef broadcast_translation(c, in_avals, out_avals, in_vals, *, shape, axes):\n  del c\n  (x,), (out_aval,) = in_vals, out_avals\n  dims_complement = [i for i in range(len(shape)) if i not in axes]\n  return [hlo.broadcast_in_dim(aval_to_ir_type(out_aval), x, dims_complement)]\nhlo_translations[broadcast_p] = broadcast_translation\n```\n\nExample:\n```text\n@jit\ndef f(x, y):\n  print('tracing!')\n  return sin(x) * cos(y)\n```\n\nExample:\n```text\nz = f(3., 4.)  # 'tracing!' prints the first time\nprint(z)\n```\n\nExample:\n```text\ntracing!\n-0.09224219304455371\n```\n\nExample:\n```text\nz = f(4., 5.)  # 'tracing!' doesn't print, compilation cache hit!\nprint(z)\n```\n\nExample:\n```text\n-0.21467624978306993\n```\n\nExample:\n```text\n@jit\ndef f(x):\n  return reduce_sum(x, axis=0)\n\nprint(f(np.array([1., 2., 3.])))\n```\n\nExample:\n```text\n6.0\n```\n\nExample:\n```text\ndef f(x):\n  y = sin(x) * 2.\n  z = - y + x\n  return z\n\ndef deriv(f):\n  return lambda x: jvp(f, (x,), (1.,))[1]\n\nprint(    deriv(deriv(f))(3.))\nprint(jit(deriv(deriv(f)))(3.))\n```\n\nExample:\n```text\n0.2822400161197344\n0.2822400161197344\n```\n\nExample:\n```text\ndef xla_call_jvp_rule(primals, tangents, *, jaxpr, num_consts):\n  del num_consts  # Unused\n  new_jaxpr, new_consts = jvp_jaxpr(jaxpr)\n  outs = bind(xla_call_p, *new_consts, *primals, *tangents, jaxpr=new_jaxpr,\n              num_consts=len(new_consts))\n  n = len(outs) // 2\n  primals_out, tangents_out = outs[:n], outs[n:]\n  return primals_out, tangents_out\njvp_rules[xla_call_p] = xla_call_jvp_rule\n\n@lru_cache\ndef jvp_jaxpr(jaxpr: Jaxpr) -> tuple[Jaxpr, list[Any]]:\n  def jvp_traceable(*primals_and_tangents):\n    n = len(primals_and_tangents) // 2\n    primals, tangents = primals_and_tangents[:n], primals_and_tangents[n:]\n    return jvp(jaxpr_as_fun(jaxpr), primals, tangents)\n\n  in_avals = [v.aval for v in jaxpr.in_binders]\n  new_jaxpr, new_consts, _ = make_jaxpr(jvp_traceable, *in_avals, *in_avals)\n  return new_jaxpr, new_consts\n```\n\nExample:\n```text\ndef xla_call_vmap_rule(axis_size, vals_in, dims_in, *, jaxpr, num_consts):\n  del num_consts  # Unused\n  new_jaxpr, new_consts = vmap_jaxpr(jaxpr, axis_size, tuple(dims_in))\n  outs = bind(xla_call_p, *new_consts, *vals_in, jaxpr=new_jaxpr,\n              num_consts=len(new_consts))\n  return outs, [0] * len(outs)\nvmap_rules[xla_call_p] = xla_call_vmap_rule\n\n@lru_cache\ndef vmap_jaxpr(jaxpr: Jaxpr, axis_size: int, bdims_in: tuple[BatchAxis, ...]\n               ) -> tuple[Jaxpr, list[Any]]:\n  vmap_traceable = vmap(jaxpr_as_fun(jaxpr), tuple(bdims_in))\n  in_avals = [unmapped_aval(axis_size, d, v.aval)\n              for v, d in zip(jaxpr.in_binders, bdims_in)]\n  new_jaxpr, new_consts, _ = make_jaxpr(vmap_traceable, *in_avals)\n  return new_jaxpr, new_consts\n\ndef unmapped_aval(axis_size: int, batch_dim: BatchAxis, aval: ShapedArray\n                  ) -> ShapedArray:\n  if batch_dim is not_mapped:\n    return aval\n  else:\n    shape = list(aval.shape)\n    shape.insert(batch_dim, axis_size)\n    return ShapedArray(tuple(shape), aval.dtype)\n```\n\nExample:\n```text\ndef xla_call_abstract_eval_rule(*in_types, jaxpr, num_consts):\n  del num_consts  # Unused\n  jaxpr_type = typecheck_jaxpr(jaxpr)\n  if not all(t1 == t2 for t1, t2 in zip(jaxpr_type.in_types, in_types)):\n    raise TypeError\n  return jaxpr_type.out_types\nabstract_eval_rules[xla_call_p] = xla_call_abstract_eval_rule\n\ndef xla_call_translation(c, in_avals, out_avals, in_vals, *, jaxpr, num_consts):\n  del num_consts, out_avals\n  # Calling jaxpr_subcomp directly would inline. We generate a Call HLO instead.\n  with ir.InsertionPoint(c.module.body):\n    @func.func(*(aval_to_ir_type(aval) for aval in in_avals))\n    def inner_xla_call(*params):\n      return jaxpr_subcomp(c, jaxpr, params)\n    c.symbol_table.insert(inner_xla_call.func_op)\n  return func.CallOp(inner_xla_call.func_op, in_vals).results\nhlo_translations[xla_call_p] = xla_call_translation\n```\n\nExample:\n```text\n@jit\ndef f(x):\n  print('tracing!')\n  y = sin(x) * 2.\n  z = - y + x\n  return z\n\nx, xdot = 3., 1.\ny, ydot = jvp(f, (x,), (xdot,))\nprint(y)\nprint(ydot)\n```\n\nExample:\n```text\ntracing!\n2.7177599838802657\n2.979984993200891\n```\n\nExample:\n```text\ny, ydot = jvp(f, (x,), (xdot,))  # 'tracing!' not printed\n```\n\nExample:\n```text\nys = vmap(f, (0,))(np.arange(3.))\nprint(ys)\n```\n\nExample:\n```text\n[ 0.         -0.68294197  0.18140515]\n```\n\nExample:\n```text\ndef handle_result(aval: ShapedArray, buf):  # noqa: F811\n  return Array(aval, buf)\n\nclass Array:\n  buf: Any\n  aval: ShapedArray\n\n  def __init__(self, aval, buf):\n    self.aval = aval\n    self.buf = buf\n\n  dtype = property(lambda self: self.aval.dtype)\n  shape = property(lambda self: self.aval.shape)\n  ndim  = property(lambda self: self.aval.ndim)\n\n  def __array__(self): return np.asarray(self.buf)\n  def __repr__(self):  return repr(np.asarray(self.buf))\n  def __str__(self):   return str(np.asarray(self.buf))\n\n  _neg = staticmethod(neg)\n  _add = staticmethod(add)\n  _radd = staticmethod(add)\n  _mul = staticmethod(mul)\n  _rmul = staticmethod(mul)\n  _gt = staticmethod(greater)\n  _lt = staticmethod(less)\ninput_handlers[Array] = lambda x: x.buf\n\njax_types.add(Array)\n```\n\nExample:\n```text\n@jit\ndef f(x):\n  y = sin(x) * 2.\n  z = - y + x\n  return z\n\nx, xdot = 3., 1.\ny, ydot = jvp(f, (x,), (xdot,))\nprint(y)\nprint(ydot)\n```\n\nExample:\n```text\ndef pprint_xla_call(names: defaultdict[Var, str], eqn: JaxprEqn) -> PPrint:\n  lhs = pp(' '.join(var_str(names, v) for v in eqn.out_binders))\n  params_without_jaxpr = {k:v for k, v in eqn.params.items() if k != 'jaxpr'}\n  rhs = (pp(eqn.primitive.name) >> pp_params(params_without_jaxpr) >>\n         pp(' '.join(names[x] if isinstance(x, Var) else str(x.val)\n                     for x in eqn.inputs)))\n  return vcat([lhs >> pp(' = ') >> rhs,\n               pp_jaxpr(eqn.params['jaxpr']).indent(2)])\npp_rules[xla_call_p] = pprint_xla_call\n```\n\nExample:\n```text\ny, f_lin = linearize(f, x)\ny_dot = f_lin(x_dot)\n```\n\nExample:\n```text\ny, y_dot = jvp(f, (x,), (x_dot,))\n```\n\nExample:\n```text\njvp : (a -> b) -> (UnrestrictedUse a, T a) -o (UnrestrictedUse b, T b)\n```\n\nExample:\n```text\ndef split_half(lst: list[Any]) -> tuple[list[Any], list[Any]]:\n  assert not len(lst) % 2\n  return split_list(lst, len(lst) // 2)\n\ndef merge_lists(which: list[bool], l1: list[Any], l2: list[Any]) -> list[Any]:\n  l1, l2 = iter(l1), iter(l2)\n  out = [next(l2) if b else next(l1) for b in which]\n  assert next(l1, None) is next(l2, None) is None\n  return out\n```\n\nExample:\n```text\ndef linearize_flat(f, *primals_in):\n  pvals_in = ([PartialVal.known(x) for x in primals_in] +\n              [PartialVal.unknown(vspace(get_aval(x))) for x in primals_in])\n  def f_jvp(*primals_tangents_in):\n    primals_out, tangents_out = jvp(f, *split_half(primals_tangents_in))\n    return [*primals_out, *tangents_out]\n  jaxpr, pvals_out, consts = partial_eval_flat(f_jvp, pvals_in)\n  primal_pvals, _ = split_half(pvals_out)\n  assert all(pval.is_known for pval in primal_pvals)\n  primals_out = [pval.const for pval in primal_pvals]\n  f_lin = lambda *tangents: eval_jaxpr(jaxpr, [*consts, *tangents])\n  return primals_out, f_lin\n\ndef linearize(f, *primals_in):\n  primals_in_flat, in_tree = tree_flatten(primals_in)\n  f, out_tree = flatten_fun(f, in_tree)\n  primals_out_flat, f_lin_flat = linearize_flat(f, *primals_in_flat)\n  primals_out = tree_unflatten(out_tree(), primals_out_flat)\n\n  def f_lin(*tangents_in):\n    tangents_in_flat, in_tree2 = tree_flatten(tangents_in)\n    if in_tree != in_tree2: raise TypeError\n    tangents_out_flat = f_lin_flat(*tangents_in_flat)\n    return tree_unflatten(out_tree(), tangents_out_flat)\n\n  return primals_out, f_lin\n\ndef vspace(aval: ShapedArray) -> ShapedArray:\n  return raise_to_shaped(aval)  # TODO handle integers?\n```\n\nExample:\n```text\npartial_eval : ((a1, a2) -> (b1, b2)) -> a1 -> exists r. (b1, r, (r, a2) -> b2)\n```\n\nExample:\n```text\n{ lambda a:float64[] .\n  let b:float64[] = sin a\n      c:float64[] = neg b\n  in ( c ) }\n```\n\nExample:\n```text\n{ lambda a:float64[] b:float64[] .\n  let c:float64[] = sin a\n      d:float64[] = cos a\n      e:float64[] = mul d b\n      f:float64[] = neg c\n      g:float64[] = neg e\n  in ( f, g ) }\n```\n\nExample:\n```text\n{ lambda a:float64[] .\n  let c:float64[] = sin a\n      d:float64[] = cos a\n      f:float64[] = neg c\n  in ( f, d ) }\n```\n\nExample:\n```text\n{ lambda d:float64[] b:float64[] .\n  let e:float64[] = mul d b\n      g:float64[] = neg e\n  in ( g ) }\n```\n\nExample:\n```text\nclass PartialVal(NamedTuple):\n  aval: ShapedArray\n  const: Any | None\n\n  @classmethod\n  def known(cls, val: Any):\n    return PartialVal(get_aval(val), val)\n\n  @classmethod\n  def unknown(cls, aval: ShapedArray):\n    return PartialVal(aval, None)\n\n  is_known   = property(lambda self: self.const is not None)\n  is_unknown = property(lambda self: self.const is     None)\n```\n\nExample:\n```text\ndef partial_eval_flat(f: Callable, pvals_in: list[PartialVal]\n                      ) -> tuple[Jaxpr, list[PartialVal], list[Any]]:\n  with new_main(PartialEvalTrace) as main:\n    trace = PartialEvalTrace(main)\n    tracers_in = [trace.new_arg(pval) for pval in pvals_in]\n    outs = f(*tracers_in)\n    tracers_out = [full_raise(trace, out) for out in outs]\n    pvals_out = [t.pval for t in tracers_out]\n    unk_tracers_in  = [t for t in tracers_in  if t.pval.is_unknown]\n    unk_tracers_out = [t for t in tracers_out if t.pval.is_unknown]\n    jaxpr, consts = tracers_to_jaxpr(unk_tracers_in, unk_tracers_out)\n  return jaxpr, pvals_out, consts\n```\n\nExample:\n```text\nfrom weakref import ref, ReferenceType\n\nclass LambdaBindingRecipe(NamedTuple):\n  pass\n\nclass ConstRecipe(NamedTuple):\n  val: Any\n\nclass JaxprEqnRecipe(NamedTuple):\n  prim: Primitive\n  tracers_in: list['PartialEvalTracer']\n  params: dict[str, Any]\n  avals_out: list[ShapedArray]\n  tracer_refs_out: list['ReferenceType[PartialEvalTracer]']\n\nJaxprRecipe = Union[LambdaBindingRecipe, ConstRecipe, JaxprEqnRecipe]\n```\n\nExample:\n```text\nclass PartialEvalTracer(Tracer):\n  pval: PartialVal\n  recipe: JaxprRecipe | None\n\n  def __init__(self, trace, pval, recipe):\n    self._trace = trace\n    self.pval = pval\n    self.recipe = recipe\n\n  aval = property(lambda self: self.pval.aval)\n\n  def full_lower(self):\n    if self.pval.is_known:\n      return full_lower(self.pval.const)\n    return self\n```\n\nExample:\n```text\nclass PartialEvalTrace(Trace):\n  def new_arg(self, pval: PartialVal) -> Any:\n    return PartialEvalTracer(self, pval, LambdaBindingRecipe())\n\n  def lift(self, val: Any) -> PartialEvalTracer:\n    return PartialEvalTracer(self, PartialVal.known(val), None)\n  pure = lift\n\n  def instantiate_const(self, tracer: PartialEvalTracer) -> PartialEvalTracer:\n    if tracer.pval.is_unknown:\n      return tracer\n    else:\n      pval = PartialVal.unknown(raise_to_shaped(tracer.aval))\n      return PartialEvalTracer(self, pval, ConstRecipe(tracer.pval.const))\n\n  def process_primitive(self, primitive, tracers, params):\n    if all(t.pval.is_known for t in tracers):\n      return bind(primitive, *map(full_lower, tracers), **params)\n    rule = partial_eval_rules.get(primitive)\n    if rule: return rule(self, tracers, **params)\n    tracers_in = [self.instantiate_const(t) for t in tracers]\n    avals_in = [t.aval for t in tracers_in]\n    avals_out = abstract_eval_rules[primitive](*avals_in, **params)\n    tracers_out = [PartialEvalTracer(self, PartialVal.unknown(aval), None)\n                   for aval in avals_out]\n    eqn = JaxprEqnRecipe(primitive, tracers_in, params, avals_out,\n                         map(ref, tracers_out))\n    for t in tracers_out: t.recipe = eqn\n    return tracers_out\n\npartial_eval_rules = {}\n```\n\nExample:\n```text\ndef tracers_to_jaxpr(tracers_in: list[PartialEvalTracer],\n                     tracers_out: list[PartialEvalTracer]):\n  tracer_to_var: dict[int, Var] = {id(t): Var(raise_to_shaped(t.aval))\n                                   for t in tracers_in}\n  constvar_to_val: dict[int, Any] = {}\n  constid_to_var: dict[int, Var] = {}\n  processed_eqns: set[int] = set()\n  eqns: list[JaxprEqn] = []\n  for t in toposort(tracers_out, tracer_parents):\n    if isinstance(t.recipe, LambdaBindingRecipe):\n      assert id(t) in set(map(id, tracers_in))\n    elif isinstance(t.recipe, ConstRecipe):\n      val = t.recipe.val\n      var = constid_to_var.get(id(val))\n      if var is None:\n        aval = raise_to_shaped(get_aval(val))\n        var = constid_to_var[id(val)] = Var(aval)\n        constvar_to_val[var] = val\n      tracer_to_var[id(t)] = var\n    elif isinstance(t.recipe, JaxprEqnRecipe):\n      if id(t.recipe) not in processed_eqns:\n        eqns.append(recipe_to_eqn(tracer_to_var, t.recipe))\n        processed_eqns.add(id(t.recipe))\n    else:\n      raise TypeError(t.recipe)\n\n  constvars, constvals = unzip2(constvar_to_val.items())\n  in_binders = constvars + [tracer_to_var[id(t)] for t in tracers_in]\n  out_vars = [tracer_to_var[id(t)] for t in tracers_out]\n  jaxpr = Jaxpr(in_binders, eqns, out_vars)\n  typecheck_jaxpr(jaxpr)\n  return jaxpr, constvals\n\ndef recipe_to_eqn(tracer_to_var: dict[int, Var], recipe: JaxprEqnRecipe\n                  ) -> JaxprEqn:\n  inputs = [tracer_to_var[id(t)] for t in recipe.tracers_in]\n  out_binders = [Var(aval) for aval in recipe.avals_out]\n  for t_ref, var in zip(recipe.tracer_refs_out, out_binders):\n    if t_ref() is not None: tracer_to_var[id(t_ref())] = var\n  return JaxprEqn(recipe.prim, inputs, recipe.params, out_binders)\n\ndef tracer_parents(t: PartialEvalTracer) -> list[PartialEvalTracer]:\n  return t.recipe.tracers_in if isinstance(t.recipe, JaxprEqnRecipe) else []\n```\n\nExample:\n```text\ndef toposort(out_nodes: list[Any], parents: Callable[[Any], list[Any]]):\n  if not out_nodes: return []\n  out_nodes = remove_duplicates(out_nodes)\n\n  child_counts = {}\n  stack = list(out_nodes)\n  while stack:\n    node = stack.pop()\n    if id(node) in child_counts:\n      child_counts[id(node)] += 1\n    else:\n      child_counts[id(node)] = 1\n      stack.extend(parents(node))\n  for node in out_nodes:\n    child_counts[id(node)] -= 1\n\n  sorted_nodes = []\n  childless_nodes = [node for node in out_nodes if not child_counts[id(node)]]\n  while childless_nodes:\n    node = childless_nodes.pop()\n    sorted_nodes.append(node)\n    for parent in parents(node):\n      if child_counts[id(parent)] == 1:\n        childless_nodes.append(parent)\n      else:\n        child_counts[id(parent)] -= 1\n\n  sorted_nodes = sorted_nodes[::-1]\n  check_toposort(sorted_nodes, parents)\n  return sorted_nodes\n\ndef remove_duplicates(lst):\n  seen = set()\n  return [x for x in lst if id(x) not in seen and not seen.add(id(x))]\n\ndef check_toposort(nodes: list[Any], parents: Callable[[Any], list[Any]]):\n  seen = set()\n  for node in nodes:\n    assert all(id(parent) in seen for parent in parents(node))\n    seen.add(id(node))\n```\n\nExample:\n```text\ny, sin_lin = linearize(sin, 3.)\nprint(y, sin(3.))\nprint(sin_lin(1.), cos(3.))\n```\n\nExample:\n```text\n0.1411200080598672 0.1411200080598672\n-0.9899924966004454 -0.9899924966004454\n```\n\nExample:\n```text\ndef xla_call_partial_eval(trace, tracers, *, jaxpr, num_consts):\n  del num_consts  # Unused\n  in_unknowns = [not t.pval.is_known for t in tracers]\n  jaxpr1, jaxpr2, out_unknowns, num_res = partial_eval_jaxpr(jaxpr, in_unknowns)\n  known_tracers, unknown_tracers = partition_list(in_unknowns, tracers)\n  known_vals = [t.pval.const for t in known_tracers]\n  outs1_res = bind(xla_call_p, *known_vals, jaxpr=jaxpr1, num_consts=0)\n  outs1, res = split_list(outs1_res, len(jaxpr1.outs) - num_res)\n  res_tracers = [trace.instantiate_const(full_raise(trace, x)) for x in res]\n  outs2 = [PartialEvalTracer(trace, PartialVal.unknown(v.aval), None)\n           for v in jaxpr2.outs]\n  eqn = JaxprEqnRecipe(xla_call_p, res_tracers + unknown_tracers,\n                       dict(jaxpr=jaxpr2, num_consts=0),\n                       [v.aval for v in jaxpr2.outs], map(ref, outs2))\n  for t in outs2: t.recipe = eqn\n  return merge_lists(out_unknowns, outs1, outs2)\npartial_eval_rules[xla_call_p] = xla_call_partial_eval\n\ndef partial_eval_jaxpr(jaxpr: Jaxpr, in_unknowns: list[bool],\n                       instantiate: list[bool] | None = None,\n                       ) -> tuple[Jaxpr, Jaxpr, list[bool], int]:\n  env: dict[Var, bool] = {}\n  residuals: set[Var] = set()\n\n  def read(x: Atom) -> bool:\n    return type(x) is Var and env[x]\n\n  def write(unk: bool, v: Var) -> None:\n    env[v] = unk\n\n  def new_res(x: Atom) -> Atom:\n    if type(x) is Var: residuals.add(x)\n    return x\n\n  eqns1, eqns2 = [], []\n  map(write, in_unknowns, jaxpr.in_binders)\n  for eqn in jaxpr.eqns:\n    unks_in = map(read, eqn.inputs)\n    rule = partial_eval_jaxpr_rules.get(eqn.primitive)\n    if rule:\n      eqn1, eqn2, unks_out, res = rule(unks_in, eqn)\n      eqns1.append(eqn1); eqns2.append(eqn2); residuals.update(res)\n      map(write, unks_out, eqn.out_binders)\n    elif any(unks_in):\n      inputs = [v if unk else new_res(v) for unk, v in zip(unks_in, eqn.inputs)]\n      eqns2.append(JaxprEqn(eqn.primitive, inputs, eqn.params, eqn.out_binders))\n      map(partial(write, True), eqn.out_binders)\n    else:\n      eqns1.append(eqn)\n      map(partial(write, False), eqn.out_binders)\n  out_unknowns = map(read, jaxpr.outs)\n  if instantiate is not None:\n    for v, uk, inst in zip(jaxpr.outs, out_unknowns, instantiate):\n      if inst and not uk: new_res(v)\n    out_unknowns = map(op.or_, out_unknowns, instantiate)\n\n  residuals, num_res = list(residuals), len(residuals)\n  assert all(type(v) is Var for v in residuals), residuals\n\n  ins1, ins2 = partition_list(in_unknowns, jaxpr.in_binders)\n  outs1, outs2 = partition_list(out_unknowns, jaxpr.outs)\n\n  jaxpr1 = Jaxpr(ins1, eqns1, outs1 + residuals)\n  jaxpr2 = Jaxpr(residuals + ins2, eqns2, outs2)\n  typecheck_partial_eval_jaxpr(jaxpr, in_unknowns, out_unknowns, jaxpr1, jaxpr2)\n\n  return jaxpr1, jaxpr2, out_unknowns, num_res\n\ndef typecheck_partial_eval_jaxpr(jaxpr, unks_in, unks_out, jaxpr1, jaxpr2):\n  jaxprty = typecheck_jaxpr(jaxpr)    # (a1,  a2) -> (b1, b2 )\n  jaxpr1ty = typecheck_jaxpr(jaxpr1)  #  a1       -> (b1, res)\n  jaxpr2ty = typecheck_jaxpr(jaxpr2)  # (res, a2) -> b2\n\n  a1, a2 = partition_list(unks_in, jaxprty.in_types)\n  b1, b2 = partition_list(unks_out, jaxprty.out_types)\n  b1_, res = split_list(jaxpr1ty.out_types, len(b1))\n  res_, a2_ = split_list(jaxpr2ty.in_types, len(res))\n  b2_ = jaxpr2ty.out_types\n\n  if jaxpr1ty.in_types != a1: raise TypeError\n  if jaxpr2ty.out_types != b2: raise TypeError\n  if b1 != b1_: raise TypeError\n  if res != res_: raise TypeError\n  if a2 != a2_: raise TypeError\n  if b2 != b2_: raise TypeError\n\npartial_eval_jaxpr_rules = {}\n\ndef xla_call_peval_eqn(unks_in: list[bool], eqn: JaxprEqn,\n                       ) -> tuple[JaxprEqn, JaxprEqn, list[bool], list[Var]]:\n  jaxpr = eqn.params['jaxpr']\n  jaxpr1, jaxpr2, unks_out, num_res = partial_eval_jaxpr(jaxpr, unks_in)\n  ins1, ins2 = partition_list(unks_in, eqn.inputs)\n  out_binders1, out_binders2 = partition_list(unks_out, eqn.out_binders)\n  residuals = [Var(v.aval) for v in jaxpr2.in_binders[:num_res]]\n  eqn1 = JaxprEqn(xla_call_p, ins1, dict(jaxpr=jaxpr1, num_consts=0),\n                  out_binders1 + residuals)\n  eqn2 = JaxprEqn(xla_call_p, residuals + ins2,\n                  dict(jaxpr=jaxpr2, num_consts=0), out_binders2)\n  return eqn1, eqn2, unks_out, residuals\npartial_eval_jaxpr_rules[xla_call_p] = xla_call_peval_eqn\n```\n\nExample:\n```text\n@jit\ndef f(x):\n  y = sin(x) * 2.\n  z = - y + x\n  return z\n\ny, f_lin = linearize(f, 3.)\ny_dot = f_lin(1.)\nprint(y, y_dot)\n```\n\nExample:\n```text\n@jit\ndef f(x):\n  y = sin(x) * 2.\n  z = g(x, y)\n  return z\n\n@jit\ndef g(x, y):\n  return cos(x) + y\n\ny, f_lin = linearize(f, 3.)\ny_dot = f_lin(1.)\nprint(y, y_dot)\n```\n\nExample:\n```text\n-0.7077524804807109 -2.121105001260758\n```\n\nExample:\n```text\nlinearize : (a -> b) -> a -> (b, T a -o T b)\nvjp       : (a -> b) -> a -> (b, T b -o T a)\n```\n\nExample:\n```text\ndef vjp(f, x):\n  y, f_lin = linearize(f, x)\n  f_vjp = lambda y_bar: transpose(f_lin)(y_bar)\n  return y, f_vjp\n```\n\nExample:\n```text\ndef vjp_flat(f, *primals_in):\n  pvals_in = ([PartialVal.known(x) for x in primals_in] +\n              [PartialVal.unknown(vspace(get_aval(x))) for x in primals_in])\n  primal_pvals_in, tangent_pvals_in = split_half(pvals_in)\n  def f_jvp(*primals_tangents_in):\n    primals_out, tangents_out = jvp(f, *split_half(primals_tangents_in))\n    return [*primals_out, *tangents_out]\n  jaxpr, pvals_out, consts = partial_eval_flat(f_jvp, pvals_in)  # linearize\n  primal_pvals, _ = split_half(pvals_out)\n  assert all(pval.is_known for pval in primal_pvals)\n  primals_out = [pval.const for pval in primal_pvals]\n  transpose_inputs = consts + [UndefPrimal(p.aval) for p in tangent_pvals_in]\n  f_vjp = lambda *cts: eval_jaxpr_transposed(jaxpr, transpose_inputs, cts)\n  return primals_out, f_vjp\n\ndef vjp(f, *primals_in):\n  primals_in_flat, in_tree = tree_flatten(primals_in)\n  f, out_tree = flatten_fun(f, in_tree)\n  primals_out_flat, f_vjp_flat = vjp_flat(f, *primals_in_flat)\n  primals_out = tree_unflatten(out_tree(), primals_out_flat)\n\n  def f_vjp(*cotangents_out):\n    cotangents_out_flat, _ = tree_flatten(cotangents_out)\n    cotangents_in_flat = f_vjp_flat(*cotangents_out_flat)\n    return tree_unflatten(in_tree, cotangents_in_flat)\n\n  return primals_out, f_vjp\n\nclass UndefPrimal(NamedTuple):\n  aval: ShapedArray\n\nregister_pytree_node(UndefPrimal,\n                     lambda u: (u.aval, ()),\n                     lambda aval, _: UndefPrimal(aval))\n```\n\nExample:\n```text\n# NB: the analogous function in JAX is called 'backward_pass'\ndef eval_jaxpr_transposed(jaxpr: Jaxpr, args: list[Any], cotangents: list[Any]\n                          ) -> list[Any]:\n  primal_env: dict[Var, Any] = {}\n  ct_env: dict[Var, Any] = {}\n\n  def read_primal(x: Atom) -> Any:\n    return primal_env.get(x, UndefPrimal(x.aval)) if type(x) is Var else x.val\n\n  def write_primal(v: Var, val: Any) -> None:\n    if type(val) is not UndefPrimal:\n      primal_env[v] = val\n\n  def read_cotangent(v: Var) -> Any:\n    return ct_env.pop(v, np.zeros(v.aval.shape, v.aval.dtype))\n\n  def write_cotangent(x: Atom, val: Any):\n    if type(x) is Var and val is not None:\n      ct_env[x] = add(ct_env[x], val) if x in ct_env else val\n\n  map(write_primal, jaxpr.in_binders, args)\n  map(write_cotangent, jaxpr.outs, cotangents)\n  for eqn in jaxpr.eqns[::-1]:\n    primals_in = map(read_primal, eqn.inputs)\n    cts_in = map(read_cotangent, eqn.out_binders)\n    rule = transpose_rules[eqn.primitive]\n    cts_out = rule(cts_in, *primals_in, **eqn.params)\n    map(write_cotangent, eqn.inputs, cts_out)\n\n  return [read_cotangent(v) for v, x in zip(jaxpr.in_binders, args)\n          if type(x) is UndefPrimal]\n\ntranspose_rules = {}\n```\n\nExample:\n```text\ndef mul_transpose_rule(cts, x, y):\n  z_bar, = cts\n  assert (type(x) is UndefPrimal) ^ (type(y) is UndefPrimal)\n  return [mul(z_bar, y), None] if type(x) is UndefPrimal else [None, mul(x, z_bar)]\ntranspose_rules[mul_p] = mul_transpose_rule\n\ndef neg_transpose_rule(cts, x):\n  ybar, = cts\n  assert type(x) is UndefPrimal\n  return [neg(ybar)]\ntranspose_rules[neg_p] = neg_transpose_rule\n\ndef add_transpose_rule(cts, x, y):\n  z_bar, = cts\n  return [z_bar, z_bar]\ntranspose_rules[add_p] = add_transpose_rule\n\ndef reduce_sum_transpose_rule(cts, x, *, axis):\n  y_bar, = cts\n  return [broadcast(y_bar, x.aval.shape, axis)]\ntranspose_rules[reduce_sum_p] = reduce_sum_transpose_rule\n\ndef xla_call_transpose_rule(cts, *invals, jaxpr, num_consts):\n  del num_consts  # Unused\n  undef_primals = [type(x) is UndefPrimal for x in invals]\n  transposed_jaxpr, new_consts = transpose_jaxpr(jaxpr, tuple(undef_primals))\n  residuals, _ = partition_list(undef_primals, invals)\n  outs = bind(xla_call_p, *new_consts, *residuals, *cts,\n              jaxpr=transposed_jaxpr, num_consts=len(new_consts))\n  outs = iter(outs)\n  return [next(outs) if undef else None for undef in undef_primals]\ntranspose_rules[xla_call_p] = xla_call_transpose_rule\n\n@lru_cache\ndef transpose_jaxpr(jaxpr: Jaxpr, undef_primals: tuple[bool, ...]\n                    ) -> tuple[Jaxpr, list[Any]]:\n  avals_in, avals_out = typecheck_jaxpr(jaxpr)\n  traceable = partial(eval_jaxpr_transposed, jaxpr)\n  args = [UndefPrimal(a) if u else a for a, u in zip(avals_in, undef_primals)]\n  trans_jaxpr, consts, _ = make_jaxpr(traceable, tuple(args), tuple(avals_out))\n  typecheck_jaxpr(trans_jaxpr)\n  return trans_jaxpr, consts\n```\n\nExample:\n```text\ndef grad(f):\n  def gradfun(x, *xs):\n    y, f_vjp = vjp(f, x, *xs)\n    if np.shape(y) != (): raise TypeError\n    x_bar, *_ = f_vjp(np.ones(np.shape(y), np.result_type(y)))\n    return x_bar\n  return gradfun\n```\n\nExample:\n```text\ny, f_vjp = vjp(sin, 3.)\nprint(f_vjp(1.), cos(3.))\n```\n\nExample:\n```text\n(np.float64(-0.9899924966004454),) -0.9899924966004454\n```\n\nExample:\n```text\ndef f(x):\n  y = sin(x) * 2.\n  z = - y + x\n  return z\n\nprint(grad(f)(3.))\n```\n\nExample:\n```text\n2.979984993200891\n```\n\nExample:\n```text\n@jit\ndef f(x):\n  y = x * 2.\n  z = g(y)\n  return z\n\n@jit\ndef g(x):\n  return cos(x) * 2.\n\nprint(grad(f)(3.))\n```\n\nExample:\n```text\n1.1176619927957034\n```\n\nExample:\n```text\n# from core_test.py fun_with_nested_calls_2\ndef foo(x):\n  @jit\n  def bar(y):\n    def baz(w):\n      q = jit(lambda x: y)(x)\n      q = q + jit(lambda: y)()\n      q = q + jit(lambda y: w + y)(y)\n      q = jit(lambda w: jit(sin)(x) * y)(1.0) + q\n      return q\n    p, t = jvp(baz, (x + 1.0,), (y,))\n    return t + (x * p)\n  return bar(x)\n\ndef assert_allclose(*vals):\n  for v1, v2 in zip(vals[:-1], vals[1:]):\n    np.testing.assert_allclose(v1, v2)\n\nans1 = f(3.)\nans2 = jit(f)(3.)\nans3, _ = jvp(f, (3.,), (5.,))\nans4, _ = jvp(jit(f), (3.,), (5.,))\nassert_allclose(ans1, ans2, ans3, ans4)\n\nderiv1 = grad(f)(3.)\nderiv2 = grad(jit(f))(3.)\nderiv3 = jit(grad(jit(f)))(3.)\n_, deriv4 = jvp(f, (3.,), (1.,))\n_, deriv5 = jvp(jit(f), (3.,), (1.,))\nassert_allclose(deriv1, deriv2, deriv3, deriv4, deriv5)\n\nhess1 = grad(grad(f))(3.)\nhess2 = grad(grad(jit(f)))(3.)\nhess3 = grad(jit(grad(f)))(3.)\nhess4 = jit(grad(grad(f)))(3.)\n_, hess5 = jvp(grad(f), (3.,), (1.,))\n_, hess6 = jvp(jit(grad(f)), (3.,), (1.,))\n_, hess7 = jvp(jit(grad(f)), (3.,), (1.,))\nassert_allclose(hess1, hess2, hess3, hess4, hess5, hess6, hess7)\n```\n\nExample:\n```text\ndef cond(pred, true_fn, false_fn, *operands):\n  avals_in = [raise_to_shaped(get_aval(x)) for x in operands]\n  true_jaxpr, true_consts, out_tree = make_jaxpr(true_fn, *avals_in)\n  false_jaxpr, false_consts, out_tree_ = make_jaxpr(false_fn, *avals_in)\n  if out_tree != out_tree_: raise TypeError\n  true_jaxpr, false_jaxpr = _join_jaxpr_consts(\n      true_jaxpr, false_jaxpr, len(true_consts), len(false_consts))\n  if typecheck_jaxpr(true_jaxpr) != typecheck_jaxpr(false_jaxpr):\n    raise TypeError\n  outs = bind_cond(pred, *true_consts, *false_consts, *operands,\n                   true_jaxpr=true_jaxpr, false_jaxpr=false_jaxpr)\n  return tree_unflatten(out_tree, outs)\ncond_p = Primitive('cond')\n\ndef _join_jaxpr_consts(jaxpr1: Jaxpr, jaxpr2: Jaxpr, n1: int, n2: int\n                       ) -> tuple[Jaxpr, Jaxpr]:\n  jaxpr1_type, jaxpr2_type = typecheck_jaxpr(jaxpr1), typecheck_jaxpr(jaxpr2)\n  assert jaxpr1_type.in_types[n1:] == jaxpr2_type.in_types[n2:]\n  consts1, rest1 = split_list(jaxpr1.in_binders, n1)\n  consts2, rest2 = split_list(jaxpr2.in_binders, n2)\n  new_jaxpr1 = Jaxpr(consts1 + consts2 + rest1, jaxpr1.eqns, jaxpr1.outs)\n  new_jaxpr2 = Jaxpr(consts1 + consts2 + rest2, jaxpr2.eqns, jaxpr2.outs)\n  return new_jaxpr1, new_jaxpr2\n\ndef bind_cond(pred, *args, true_jaxpr, false_jaxpr):\n  assert len(args) == len(true_jaxpr.in_binders) == len(false_jaxpr.in_binders)\n  return bind(cond_p, pred, *args, true_jaxpr=true_jaxpr, false_jaxpr=false_jaxpr)\n```\n\nExample:\n```text\ndef cond_impl(pred, *operands, true_jaxpr, false_jaxpr):\n  if pred:\n    return eval_jaxpr(true_jaxpr, operands)\n  else:\n    return eval_jaxpr(false_jaxpr, operands)\nimpl_rules[cond_p] = cond_impl\n```\n\nExample:\n```text\nout = cond(True, lambda: 3, lambda: 4)\nprint(out)\n```\n\nExample:\n```text\n3\n```\n\nExample:\n```text\ndef cond_jvp_rule(primals, tangents, *, true_jaxpr, false_jaxpr):\n  pred, *primals = primals\n  _   , *tangents = tangents\n  true_jaxpr , true_consts  = jvp_jaxpr(true_jaxpr)\n  false_jaxpr, false_consts = jvp_jaxpr(false_jaxpr)\n  true_jaxpr, false_jaxpr = _join_jaxpr_consts(\n      true_jaxpr, false_jaxpr, len(true_consts), len(false_consts))\n  assert typecheck_jaxpr(true_jaxpr) == typecheck_jaxpr(false_jaxpr)\n  outs = bind_cond(pred, *true_consts, *false_consts, *primals, *tangents,\n                   true_jaxpr=true_jaxpr, false_jaxpr=false_jaxpr)\n  primals_out, tangents_out = split_half(outs)\n  return primals_out, tangents_out\njvp_rules[cond_p] = cond_jvp_rule\n```\n\nExample:\n```text\nout, out_tan = jvp(lambda x: cond(True, lambda: x * x, lambda: 0.), (1.,), (1.,))\nprint(out_tan)\n```\n\nExample:\n```text\n2.0\n```\n\nExample:\n```text\ndef cond_vmap_rule(axis_size, vals_in, dims_in, *, true_jaxpr, false_jaxpr):\n  pred    , *vals_in = vals_in\n  pred_dim, *dims_in = dims_in\n  if pred_dim is not not_mapped: raise NotImplementedError  # TODO\n  true_jaxpr, true_consts = vmap_jaxpr(true_jaxpr, axis_size, tuple(dims_in))\n  false_jaxpr, false_consts = vmap_jaxpr(false_jaxpr, axis_size, tuple(dims_in))\n  true_jaxpr, false_jaxpr = _join_jaxpr_consts(\n      true_jaxpr, false_jaxpr, len(true_consts), len(false_consts))\n  assert typecheck_jaxpr(true_jaxpr) == typecheck_jaxpr(false_jaxpr)\n  outs = bind_cond(pred, *true_consts, *false_consts, *vals_in,\n                   true_jaxpr=true_jaxpr, false_jaxpr=false_jaxpr)\n  return outs, [0] * len(outs)\nvmap_rules[cond_p] = cond_vmap_rule\n```\n\nExample:\n```text\nxs = np.array([1., 2., 3])\nout = vmap(lambda x: cond(True, lambda: x + 1., lambda: 0.), (0,))(xs)\nprint(out)\n```\n\nExample:\n```text\n[2. 3. 4.]\n```\n\nExample:\n```text\n{ lambda a:float32[] .\n  let\n  in ( a ) }\n```\n\nExample:\n```text\n{ lambda a:float32[] .\n  let\n  in ( 0. ) }\n```\n\nExample:\n```text\ndef cond_abstract_eval(pred_type, *in_types, true_jaxpr, false_jaxpr):\n  if pred_type != ShapedArray((), np.dtype('bool')): raise TypeError\n  jaxpr_type = typecheck_jaxpr(true_jaxpr)\n  if jaxpr_type != typecheck_jaxpr(false_jaxpr):\n    raise TypeError\n  if not all(t1 == t2 for t1, t2 in zip(jaxpr_type.in_types, in_types)):\n    raise TypeError\n  return jaxpr_type.out_types\nabstract_eval_rules[cond_p] = cond_abstract_eval\n\ndef cond_translation(c, in_avals, out_avals, in_vals, *, true_jaxpr, false_jaxpr):\n  del in_avals  # Unused\n  pred, *in_vals = in_vals\n\n  op = hlo.IfOp([aval_to_ir_type(aval) for aval in out_avals], pred)\n  with ir.InsertionPoint(op.true_branch.blocks.append()):\n    hlo.return_(jaxpr_subcomp(c, true_jaxpr, in_vals))\n  with ir.InsertionPoint(op.false_branch.blocks.append()):\n    hlo.return_(jaxpr_subcomp(c, false_jaxpr, in_vals))\n  return op.results\n\nhlo_translations[cond_p] = cond_translation\n```\n\nExample:\n```text\nout = jit(lambda: cond(False, lambda: 1, lambda: 2))()\nprint(out)\n```\n\nExample:\n```text\n2\n```\n\nExample:\n```text\ndef cond_partial_eval(trace, tracers, *, true_jaxpr, false_jaxpr):\n  pred_tracer, *tracers = tracers\n  assert pred_tracer.pval.is_known\n  pred = pred_tracer.pval.const\n  in_uks = [not t.pval.is_known for t in tracers]\n\n  *jaxprs, out_uks, num_res = _cond_partial_eval(true_jaxpr, false_jaxpr, in_uks)\n  t_jaxpr1, f_jaxpr1, t_jaxpr2, f_jaxpr2 = jaxprs\n\n  known_tracers, unknown_tracers = partition_list(in_uks, tracers)\n  known_vals = [t.pval.const for t in known_tracers]\n  outs1_res = bind_cond(pred, *known_vals,\n                        true_jaxpr=t_jaxpr1, false_jaxpr=f_jaxpr1)\n  outs1, res = split_list(outs1_res, len(outs1_res) - num_res)\n  pred_tracer_ = trace.instantiate_const(full_raise(trace, pred_tracer))\n  res_tracers = [trace.instantiate_const(full_raise(trace, x)) for x in res]\n  outs2 = [PartialEvalTracer(trace, PartialVal.unknown(v.aval), None)\n           for v in t_jaxpr2.outs]\n  eqn = JaxprEqnRecipe(cond_p, [pred_tracer_, *res_tracers, *unknown_tracers],\n                       dict(true_jaxpr=t_jaxpr2, false_jaxpr=f_jaxpr2),\n                       [v.aval for v in t_jaxpr2.outs], map(ref, outs2))\n  for t in outs2: t.recipe = eqn\n  return merge_lists(out_uks, outs1, outs2)\npartial_eval_rules[cond_p] = cond_partial_eval\n\ndef _cond_partial_eval(true_jaxpr: Jaxpr, false_jaxpr: Jaxpr, in_uks: list[bool]\n                       ) -> tuple[Jaxpr, Jaxpr, Jaxpr, Jaxpr, list[bool], int]:\n  _, _, t_out_uks, _ = partial_eval_jaxpr(true_jaxpr , in_uks)\n  _, _, f_out_uks, _ = partial_eval_jaxpr(false_jaxpr, in_uks)\n  out_uks = map(op.or_, t_out_uks, f_out_uks)\n\n  t_jaxpr1, t_jaxpr2, _, t_nres = partial_eval_jaxpr(true_jaxpr , in_uks, out_uks)\n  f_jaxpr1, f_jaxpr2, _, f_nres = partial_eval_jaxpr(false_jaxpr, in_uks, out_uks)\n\n  t_jaxpr1, f_jaxpr1 = _join_jaxpr_res(t_jaxpr1, f_jaxpr1, t_nres, f_nres)\n  t_jaxpr2, f_jaxpr2 = _join_jaxpr_consts(t_jaxpr2, f_jaxpr2, t_nres, f_nres)\n  assert typecheck_jaxpr(t_jaxpr1) == typecheck_jaxpr(f_jaxpr1)\n  assert typecheck_jaxpr(t_jaxpr2) == typecheck_jaxpr(f_jaxpr2)\n  num_res = t_nres + f_nres\n\n  return t_jaxpr1, f_jaxpr1, t_jaxpr2, f_jaxpr2, out_uks, num_res\n\ndef _join_jaxpr_res(jaxpr1: Jaxpr, jaxpr2: Jaxpr, n1: int, n2: int\n                    ) -> tuple[Jaxpr, Jaxpr]:\n  jaxpr1_type, jaxpr2_type = typecheck_jaxpr(jaxpr1), typecheck_jaxpr(jaxpr2)\n  out_types1, _ = split_list(jaxpr1_type.out_types, len(jaxpr1.outs) - n1)\n  out_types2, _ = split_list(jaxpr2_type.out_types, len(jaxpr2.outs) - n2)\n  assert out_types1 == out_types2\n  outs1, res1 = split_list(jaxpr1.outs, len(jaxpr1.outs) - n1)\n  outs2, res2 = split_list(jaxpr2.outs, len(jaxpr2.outs) - n2)\n  zeros_like1 = [Lit(np.zeros(v.aval.shape, v.aval.dtype)) for v in res1]\n  zeros_like2 = [Lit(np.zeros(v.aval.shape, v.aval.dtype)) for v in res2]\n  new_jaxpr1 = Jaxpr(jaxpr1.in_binders, jaxpr1.eqns, outs1 + res1 + zeros_like2)\n  new_jaxpr2 = Jaxpr(jaxpr2.in_binders, jaxpr2.eqns, outs2 + zeros_like1 + res2)\n  return new_jaxpr1, new_jaxpr2\n```\n\nExample:\n```text\n_, f_lin = linearize(lambda x: cond(True, lambda: x, lambda: 0.), 1.)\nout = f_lin(3.14)\nprint(out)\n```\n\nExample:\n```text\n3.14\n```\n\nExample:\n```text\ndef cond_peval_eqn(unks_in: list[bool], eqn: JaxprEqn,\n                   ) -> tuple[JaxprEqn, JaxprEqn, list[bool], list[Atom]]:\n  pred_unk, *unks_in = unks_in\n  assert not pred_unk\n  true_jaxpr, false_jaxpr = eqn.params['true_jaxpr'], eqn.params['false_jaxpr']\n  *jaxprs, unks_out, num_res = _cond_partial_eval(true_jaxpr, false_jaxpr, unks_in)\n  t_jaxpr1, f_jaxpr1, t_jaxpr2, f_jaxpr2 = jaxprs\n  ins1, ins2 = partition_list(unks_in, eqn.inputs[1:])\n  outs1, outs2 = partition_list(unks_out, eqn.out_binders)\n  residuals, _ = split_list(t_jaxpr2.in_binders, num_res)\n  eqn1 = JaxprEqn(cond_p, [eqn.inputs[0], *ins1],\n                  dict(true_jaxpr=t_jaxpr1, false_jaxpr=f_jaxpr1),\n                  outs1 + residuals)\n  eqn2 = JaxprEqn(cond_p, [eqn.inputs[0], *residuals, *ins2],\n                  dict(true_jaxpr=t_jaxpr2, false_jaxpr=f_jaxpr2),\n                  outs2)\n  res = [eqn.inputs[0], *residuals] if type(eqn.inputs[0]) is Var else residuals\n  return eqn1, eqn2, unks_out, res\npartial_eval_jaxpr_rules[cond_p] = cond_peval_eqn\n```\n\nExample:\n```text\n_, f_lin = linearize(jit(lambda x: cond(True, lambda: x, lambda: 0.)), 1.)\nout = f_lin(3.14)\nprint(out)\n```\n\nExample:\n```text\ndef cond_transpose_rule(cts, pred, *invals, true_jaxpr, false_jaxpr):\n  undef_primals = tuple(type(x) is UndefPrimal for x in invals)\n  true_jaxpr, true_consts = transpose_jaxpr(true_jaxpr, undef_primals)\n  false_jaxpr, false_consts = transpose_jaxpr(false_jaxpr, undef_primals)\n  true_jaxpr, false_jaxpr = _join_jaxpr_consts(\n      true_jaxpr, false_jaxpr, len(true_consts), len(false_consts))\n  res = [x for x in invals if type(x) is not UndefPrimal]\n  outs = bind_cond(pred, *true_consts, *false_consts, *res, *cts,\n                   true_jaxpr=true_jaxpr, false_jaxpr=false_jaxpr)\n  outs = iter(outs)\n  return [None] + [next(outs) if type(x) is UndefPrimal else None for x in invals]\ntranspose_rules[cond_p] = cond_transpose_rule\n```\n\nExample:\n```text\nout = grad(lambda x: cond(True, lambda: x * x, lambda: 0.))(1.)\nprint(out)\n```\n\nExample:\n```text\ndef pprint_cond(names: defaultdict[Var, str], eqn: JaxprEqn) -> PPrint:\n  true_jaxpr, false_jaxpr = eqn.params['true_jaxpr'], eqn.params['false_jaxpr']\n  new_params = {k:v for k, v in eqn.params.items() if not k.endswith('jaxpr')}\n  lhs = pp(' '.join(var_str(names, v) for v in eqn.out_binders))\n  rhs = (pp(eqn.primitive.name) >> pp_params(new_params) >>\n         pp(' '.join(names[x] if isinstance(x, Var) else str(x.val)\n                     for x in eqn.inputs)))\n  return vcat([lhs >> pp(' = ') >> rhs,\n               pp_jaxpr(true_jaxpr).indent(2),\n               pp_jaxpr(false_jaxpr).indent(2)])\npp_rules[cond_p] = pprint_cond\n```\n\n<|endofdoc|>","metadata":{"transformedAt":"2026-08-18T15:17:14.886Z","totalSectionsIncluded":0,"totalCodeBlocksIncluded":140,"totalLines":2616,"estimatedTokens":18694}}140