A simplified and automated orchestration workflow to perform ML end-to-end (E2E) model tests and benchmarking on Cloud VMs across different frameworks.

Tokamax: A GPU and TPU kernel library.

A simple, performant and scalable Jax LLM!

xpk (Accelerated Processing Kit, pronounced x-p-k,) is a software tool to help Cloud developers to orchestrate training jobs on accelerators such as TPUs and GPUs on GKE.

An Extensible Deep Learning Library

Testing framework for Deep Learning models (Tensorflow and PyTorch) on Google Cloud hardware accelerators (TPU and GPU)

🤗 Diffusers: State-of-the-art diffusion models for image and audio generation in PyTorch

Enabling PyTorch on Google TPU

TorchBench is a collection of open source benchmarks used to evaluate PyTorch performance.

A Python SDK for Vertex AI, a fully managed, end-to-end platform for data science and machine learning.

🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.

Flax is a neural network library for JAX that is designed for flexibility.

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more