Using JAX to accelerate our research
DeepMind engineers accelerate our research by building tools, scaling up algorithms, and creating challenging virtual and physical worlds for training and testing artificial intelligence (AI) systems. As part of this work, we constantly evaluate new machine learning libraries and frameworks.
Using JAX to accelerate our research — Google DeepMind Skip to main content December 4, 2020 Research Using JAX to accelerate our research David Budden, Matteo Hessel Share DeepMind engineers accelerate our research by building tools, scaling up algorithms, and creating challenging virtual and physical worlds for training and testing artificial intelligence (AI) systems. As part of this work, we constantly evaluate new machine learning libraries and frameworks. Recently, we've found that an increasing number of projects are well served by JAX , a machine learning framework developed by Google
related reading
- Why You Should (or Shouldn't) be Using Google's JAX in 2023assemblyai.com
- GitHub - jax-ml/jax: Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and moregithub.com
- How to think in JAX — JAX documentationdocs.jax.dev
- Build a Transformer in JAX from scratch: how to write and train your own models | AI Summertheaisummer.com
- Why You Should (or Shouldn't) be Using Google's JAX in 2023assemblyai.com
- GitHub - google-research/tuning_playbook: A playbook for systematically maximizing the performance of deep learning models.github.com
- AlphaEvolve: A Gemini-powered coding agent for designing advanced algorithms — Google DeepMinddeepmind.google
- Quickstart: How to think in JAX — JAX documentationjax.readthedocs.io
- Differentiable Programming from Scratchthenumb.at
- A guide to JAX for PyTorch developers | Google Cloud Blogcloud.google.com
- GitHub - marin-community/haliax: Named Tensors for Legible Deep Learning in JAX · GitHubgithub.com
- GitHub - google/hackable_diffusiongithub.com