Easy to use and blazing fast JAX-based library for high-performance 2D/3D Discrete Element Method (DEM) simulations.

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

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

JAX bindings for the NVIDIA cuDecomp library

katmer is a powerful library for optimizing the design of optical thin films using automatic differentiation via JAX and Equinox, enabling efficient and accurate inverse design solutions.

PyTorch distributed training acceleration framework

High-performance TensorFlow library for quantitative finance.

JAX Custom Operations with C++ and CUDA (using Pybind11)

Fast and easy distributed model training examples.

Dataloading for JAX

Extending JAX with custom C++ and CUDA code

Modern Graph TensorFlow implementation of Super-Resolution GAN

ALBERT model Pretraining and Fine Tuning using TF2.0

Provides code to serialize the different models involved in Stable Diffusion as SavedModels and to compile them with XLA.

As the quality of large language models increases, so do our expectations of what they can do. Since the release of OpenAI's GPT-2, text generation capabilities have received attention. And for good reason - these models can be used for summarization, translation, and even real-time learning in some language tasks.

Versatile Data Ingestion Pipelines for Jax

A Multivariate Gaussian Bayes classifier written using JAX

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

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