# jax

**Repository Path**: mirrors_deepmind/jax

## Basic Information

- **Project Name**: jax
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2021-02-27
- **Last Updated**: 2025-09-28

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# DeepMind JAX Ecosystem

The JAX Ecosystem is a thriving community of
[JAX](https://github.com/google/jax) libraries. Here are some libraries
developed at DeepMind:

<table>
  <thead>
    <tr>
      <th><strong>Topic</strong></th>
      <th><strong>Library</strong></th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <th>Neural networks</th>
      <td><a href="https://github.com/deepmind/dm-haiku">Haiku</a></td>
    </tr>
    <tr>
      <th>Optimizers</th>
      <td><a href="https://github.com/deepmind/optax">Optax</a></td>
    </tr>
    <tr>
      <th>Probability distributions</th>
      <td><a href="https://github.com/deepmind/distrax">Distrax</a></td>
    </tr>
    <tr>
      <th>Image processing</th>
      <td><a href="https://github.com/deepmind/dm_pix">PIX</a></td>
    </tr>
    <tr>
      <th>GraphNets</th>
      <td><a href="https://github.com/deepmind/jraph">Jraph</a></td>
    </tr>
    <tr>
      <th>Protein folding</th>
      <td><a href="https://github.com/deepmind/alphafold">AlphaFold</a></td>
    </tr>
    <tr>
      <th>Reinforcement learning</th>
      <td><a href="https://github.com/deepmind/rlax">RLax</a></td>
    </tr>
    <tr>
      <th>AlphaZero-like Tree Search</th>
      <td><a href="https://github.com/deepmind/mctx">Mctx</a></td>
    </tr>
    <tr>
      <th>Testing & debugging</th>
      <td><a href="https://github.com/deepmind/chex">Chex</a></td>
    </tr>
    <tr>
      <th>Second-Order Optimization</th>
      <td><a href="https://github.com/deepmind/kfac-jax">KFAC-JAX</a></td>
    </tr>
    <tr>
      <th>Mixed Precision Training</th>
      <td><a href="https://github.com/deepmind/jmp">JMP</a></td>
    </tr>
    <tr>
      <th>Audio processing</th>
      <td><a href="https://github.com/deepmind/aux">AUX</a></td>
    </tr>
    <tr>
      <th>TensorFlow to JAX conversion</th>
      <td><a href="https://github.com/deepmind/tf2jax">tf2jax</a></td>
    </tr>
    <tr>
      <th>Loopy Belief Propagation</th>
      <td><a href="https://github.com/deepmind/pgmax">PGMax</a></td>
    </tr>
    <tr>
      <th>Structured Probability Distributions</th>
      <td><a href="https://github.com/deepmind/synjax">SynJax</a></td>
    </tr>
  </tbody>
</table>

Originally, DeepMind introduced [the JAX Ecosystem in a blog post](https://www.deepmind.com/blog/using-jax-to-accelerate-our-research).