JAX

JAX is a Python library that provides accelerated matrix operations that are relevant for machine learning.

Applications

The following players have disclosed their use of JAX:

• Anthropic uses JAX for their inferencing stack.¹
• xAI used JAX for training Grok-1.²

Footnotes

1. https://www.anthropic.com/engineering/a-postmortem-of-three-recent-issues ↩

2. https://github.com/xai-org/grok-1 ↩