Jax

jax.dev

Actively hiring

As featured in

docs.jax.dev

docs.jaxstack.ai

Overview

News

Technologies

Salaries

Products

People

Growth

Financials

Overview

JAX: High performance array computing. JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.

News

News docs.jax.dev 1 month ago

Ref: mutable arrays for data plumbing and memory control - JAX documentation Read more

Report

News docs.jax.dev 1 month ago

Defining new JAX types with hijax - JAX documentation Read more

Report

News docs.jax.dev 10 months ago

Collective matrix multiplication - JAX documentation Read more

Report

News docs.jax.dev 10 months ago

Writing high-performance matrix multiplication kernels for Blackwell - JAX documentation Read more

Report

News docs.jaxstack.ai

Getting started with JAX for ML - JAX AI Stack

If you're looking to use JAX to train neural networks, check out the JAX AIStack !

Report