| Name | Version | Summary | date |
| jax-fdm |
0.8.6 |
Auto-differentiable and hardware-accelerated force density method |
2024-10-30 19:25:13 |
| jax-fem |
0.0.8 |
A GPU accelerated Finite element analysis package in JAX. |
2024-10-30 02:46:50 |
| jax-flash-attn2 |
0.0.1 |
Flash Attention Implementation with Multiple Backend Support and Sharding This module provides a flexible implementation of Flash Attention with support for different backends (GPU, TPU, CPU) and platforms (Triton, Pallas, JAX). |
2024-10-23 22:37:19 |
| exponax |
0.1.0 |
Efficient differentiable PDE solvers in JAX. |
2024-10-23 07:30:54 |
| diffusers |
0.31.0 |
State-of-the-art diffusion in PyTorch and JAX. |
2024-10-22 14:13:23 |
| flowjax |
16.0.0 |
Easy to use distributions, bijections and normalizing flows in JAX. |
2024-10-17 11:00:06 |
| trainax |
0.0.2 |
Training methodologies for Autoregressive Neural Emulators in JAX built on top of Equinox & Optax. |
2024-10-17 07:29:40 |
| pdequinox |
0.1.2 |
Neural PDE Emulator Architectures in JAX built on top of Equinox. |
2024-10-17 07:02:31 |
| jax-tqdm |
0.3.1 |
Tqdm progress bar for JAX scans and loops |
2024-10-14 23:39:39 |
| xminigrid |
0.9.1 |
JAX-accelerated meta-reinforcement learning environments inspired by XLand and MiniGrid |
2024-10-11 05:28:05 |
| GradientTransformation |
1.0.0 |
A JAX-based L-BFGS optimizer |
2024-10-09 17:42:56 |
| jax-tools |
1.0.161 |
Jax common tools library |
2024-10-09 10:23:59 |
| jax-metal |
0.1.1 |
JAX acceleration for Mac GPUs. |
2024-10-08 16:57:30 |
| thermox |
0.0.3 |
Exact OU processes with JAX |
2024-09-24 09:57:02 |
| jwave |
0.2.1 |
Fast and differentiable acoustic simulations in JAX. |
2024-09-17 11:07:15 |
| jaxdf |
0.2.8 |
A JAX-based research framework for writing differentiable numerical simulators with arbitrary discretizations |
2024-09-17 10:37:45 |
| scikit-jax |
0.0.2 |
Classical machine learning algorithms on the GPU/TPU. |
2024-09-15 16:44:39 |
| distreqx |
0.0.1 |
Lightweight JAX library of probability distributions and bijectors. |
2024-09-15 06:13:45 |
| inox |
0.6.3 |
Stainless neural networks in JAX |
2024-09-02 10:25:38 |
| jax-relax |
0.2.8 |
JAX-based Recourse Explanation Library |
2024-08-29 02:26:43 |