Any model. Any hardware. Zero compromise. Built with @ziglang / @openxla / MLIR / @bazelbuild

Multi-dimensional arrays (tensors) and numerical definitions for Elixir

Enabling PyTorch on XLA Devices (e.g. Google TPU)

JAX - A curated list of resources https://github.com/google/jax

GoMLX: An Accelerated Machine Learning Framework For Go

The purpose of this repo is to make it easy to get started with JAX, Flax, and Haiku. It contains my "Machine Learning with JAX" series of tutorials (YouTube videos and Jupyter Notebooks) as well as the content I found useful while learning about the JAX ecosystem.

Elegant and Performant Deep Learning

Felafax is building AI infra for non-NVIDIA GPUs

Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python ⚡

Extending JAX with custom C++ and CUDA code

Julia on TPUs

ALBERT model Pretraining and Fine Tuning using TF2.0

Simple and efficient RevNet-Library for PyTorch with XLA and DeepSpeed support and parameter offload

S + Autograd + XLA :: S-parameter based frequency domain circuit simulations and optimizations using JAX.

TensorFlow wheels built for latest CUDA/CuDNN and enabled performance flags: SSE, AVX, FMA; XLA

High-performance, differentiable quantum state-vector & tensor network simulator in 100% pure JAX (no classical framework overhead). Accelerated on NVIDIA GPUs and Google Cloud TPU v6e-64/v5e VM clusters up to 37 qubits! Supported by Google's TPU Research Cloud (TRC) program.

Accelerated array computing and code transformations for R

基于tensorflow1.x的预训练模型调用，支持单机多卡、梯度累积，XLA加速，混合精度。可灵活训练、验证、预测。

JAX reimplementation and bindings for the NVIDIA cuDecomp library

PyTorch distributed training acceleration framework