# google/trax

Trax — Deep Learning with Clear Code and Speed

Repository: https://github.com/google/trax
Canonical: https://ross.abutalabs.com/products/trax
Language: Python
License: Apache-2.0
License Family: permissive
Topics: jax, numpy, deep-learning, deep-reinforcement-learning, machine-learning, transformer, reinforcement-learning
Archived: true
Last push: 2025-09-26T14:37:32+00:00

## Health v2 (maintenance only)
Score: 10/100 (v2, computed 2026-09-02T17:46:02.011165+00:00)
- activity 44, release rhythm 8, longevity 100
- inputs: {"age_days": 2524, "days_push": 341, "days_rel": null, "gap_med": null, "n_releases_24m": 0}
- flags: archived
- formula: round(0.45*activity + 0.35*rhythm + 0.20*longevity); archived -> min(score, 10)

## Adoption (not part of the score)
Stars 8306, forks 820 (observed 2026-08-28T04:10:20.235855+00:00)

## What it is
Trax is an end-to-end deep learning library built on JAX and TensorFlow that focuses on clear code and speed, developed and maintained by the Google Brain team. It provides implementations of models like Transformer and Reformer, plus pre-trained models and training utilities.

## Use cases
- train transformer models for machine translation
- run a pre-trained transformer translator in a few lines of code
- implement named entity recognition with Reformer
- build deep n-gram language models
- experiment with deep learning research in JAX
- train reinforcement learning agents

## When to choose
- you want a readable, research-friendly deep learning library built on JAX
- you need ready-made Transformer or Reformer implementations with pre-trained weights
- you're learning deep learning through clear, well-documented code and notebooks

## When to avoid
- you need a production ecosystem with broad deployment support like PyTorch or TensorFlow
- you require the latest model architectures or active community development
- you want fine-grained low-level control over JAX internals

## Facets
- artifact type: library
- maturity: maintenance
- function: deep-learning, machine-learning, llm-training, reinforcement-learning, nlp
- domain: deep-learning, machine-learning, reinforcement-learning
- platform: python
- tags: jax, transformer, reformer, numpy, google-brain, research, natural-language-processing, linux, macos, gpu

## Member repositories
- google/trax (main) score 10

## Provenance
- Observed fields: from GitHub, fetched 2026-08-28T04:10:20.235855+00:00.
- Health v2: computed from the inputs above; adoption is never an input.
- Inferred fields (summary, facets, guidance): AI-extracted, prompt v1, taxonomy v1, on 2026-08-29T17:28:04.257020+00:00, confidence not recorded.
  - readme: https://github.com/google/trax (fetched 2026-08-28T04:10:20.235855+00:00, sha 4b8a169101ab)
  - registry_pypi: https://pypi.org/pypi/trax/json (fetched 2026-08-29T08:27:42.715942+00:00, sha 0b5cb270d131)
- Data as of 2026-08-30T08:39:29.467469+00:00.