273  jax(1) - Kailash Tool Reference

273.1 jax

jax - JAX — composable transformations of Python+NumPy programs.

Layer OPS
Category o-05-ai-engineering-sdks-dev-env
Safety class safe
Licence Apache-2.0
Upstream github.com/google/jax

273.1.1 NAME

jax - JAX — composable transformations of Python+NumPy programs.

273.1.2 SYNOPSIS & usage

Derived from the upstream README at generation time; verify against --help where the tool is installed.

import jax
import jax.numpy as jnp

def tanh(x):
  y = jnp.exp(-2.0 * x)
  return (1.0 - y) / (1.0 + y)

grad_tanh = jax.grad(tanh)
print(grad_tanh(1.0))
# prints 0.4199743
print(jax.grad(jax.grad(jax.grad(tanh)))(1.0))
# prints 0.62162673

273.1.3 DESCRIPTION (upstream)

JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning. JAX can automatically differentiate native Python and NumPy functions. It can differentiate through loops, branches, recursion, and closures, and it can take derivatives of derivatives of derivatives. It supports reverse-mode differentiation (a.k.a. backpropagation) via jax.grad as well as forward-mode differentiation, and the two can be composed arbitrarily to any order. JAX uses XLA to compile and scale your NumPy programs on TPUs, GPUs, and

273.1.4 CONFIGURATION (detected)

Config-surface markers detected in the upstream README (2 matches, e.g. ci-build.yaml, ci-build.yaml). A full config reference page for this tool is queued under the config-page policy — this line is the detection record, not the page.

273.1.5 SEE ALSO

ref: pin - verified 2026-10-14