spb/tools-spboucher-ai

1{2  "slug": "jax",3  "name": "JAX",4  "aliases": ["jax-ml/jax", "google/jax"],5  "vendor": { "name": "Google (jax-ml)", "slug": "google", "website": "https://docs.jax.dev" },6  "homepage_url": "https://docs.jax.dev/en/latest/",7  "docs_url": "https://docs.jax.dev/en/latest/",8  "repo_url": "https://github.com/jax-ml/jax",9  "license_url": "https://github.com/jax-ml/jax/blob/main/LICENSE",10  "product_type": "library",11  "product_status": "active",12  "summary": "Python library for NumPy-style array computing with composable transformations (jit, grad, vmap, shard_map) compiled through XLA to CPU, GPU and TPU",13  "description": "JAX is an Apache-2.0 Python library from Google for high-performance numerical computing and machine learning research. It exposes a NumPy-compatible API and a set of composable function transformations: `jax.jit` compiles functions through XLA, `jax.grad` differentiates them, `jax.vmap` vectorizes over a batch dimension and `shard_map` expresses per-device programs for multi-device execution. The same code runs on CPU, GPU and TPU backends, which is why JAX underlies much of Google's large-scale training and a growing number of academic codebases. The repository moved from google/jax to jax-ml/jax.\n\nJAX itself is deliberately minimal; neural-network layers come from Flax, Equinox or Keras, optimizers from Optax and checkpointing from Orbax. Its functional style (pure functions, explicit random keys, immutable arrays) rewards careful program design and makes transformations predictable, but it is a different mental model from PyTorch's eager, object-oriented approach. Practical limits include longer compile times for large programs, a smaller ecosystem of pretrained models than PyTorch, and GPU support that depends on CUDA builds while Apple Silicon acceleration is experimental.",14  "capabilities": ["NumPy-compatible array API", "Composable jit, grad, vmap and shard_map transformations", "XLA compilation for CPU, GPU and TPU", "Multi-device and multi-host parallelism", "Ecosystem: Flax, Equinox, Optax, Orbax"],15  "workflows": ["Write a differentiable simulation and vectorize it over parameters", "Train a Flax model on TPUs with shard_map", "Port a NumPy algorithm to GPU with minimal changes"],16  "advantages": ["Clean functional transformations", "First-class TPU support", "Apache-2.0 license"],17  "limitations": ["Compile latency for large programs", "Smaller pretrained-model ecosystem than PyTorch", "Apple Silicon GPU support is experimental"],18  "prerequisites": ["Python 3.10+; CUDA for NVIDIA GPUs or TPU access"],19  "audiences": ["ML and scientific-computing researchers", "Teams training on TPUs"],20  "alternatives": ["pytorch", "numpy", "mlx"],21  "categories": [{ "slug": "model-training", "primary": true }, { "slug": "math-stats-tools", "label": "Differentiable numerical computing" }],22  "tags": ["open-source", "python", "deep-learning", "autodiff", "tpu"],23  "use_cases": ["fine-tune-models"],24  "operating_systems": ["linux", "macos", "windows"],25  "execution_modes": ["local", "self-hosted", "cloud"],26  "interfaces": ["library"],27  "deployment_models": ["embedded"],28  "integrations": ["Flax", "Optax", "Orbax", "Equinox", "Keras", "XLA", "NumPy"],29  "data_formats": ["NumPy arrays", "Orbax checkpoints"],30  "languages": ["Python"],31  "source_model": "open-source",32  "license_spdx": "Apache-2.0",33  "pricing_model": ["open-source", "free"],34  "self_hostable": true,35  "sources": [36    { "url": "https://docs.jax.dev/en/latest/", "type": "official_site", "title": "JAX documentation", "supports": ["name", "summary", "capabilities", "integrations"], "checked_at": "2026-11-02", "note": "jax.readthedocs.io redirects here" },37    { "url": "https://github.com/jax-ml/jax", "type": "repository", "supports": ["license_spdx", "languages"], "checked_at": "2026-11-02", "note": "github.com/google/jax redirects here" },38    { "url": "https://github.com/jax-ml/jax/blob/main/LICENSE", "type": "license_page", "supports": ["license_spdx", "source_model"], "checked_at": "2026-11-02" }39  ],40  "confidence": "high",41  "discovery_method": "editorial-research",42  "depth": "standard",43  "editorial_status": "reviewed",44  "visibility": "public",45  "featured": false,46  "author": "batch-A"47}48