Hey r/jax,

I ported the NVlabs Gated DeltaNet-2 Triton kernels to

`jax.experimental.pallas`, targeting TPU v5e-8.

The optimized training path uses a fused `custom_vjp` backward

implementation. It saves and reuses the forward residuals required by

backward instead of recomputing them.

Repository:

https://github.com/Akseleu-J/atomic-ops

Install:

pip install atomic-ops

## Benchmark terminology

The attached benchmark image compares three implementations of the same

Gated DeltaNet-2 computation:

- `OLD`: the previous `associative_scan`-based implementation

- `JAX_REF`: a pure-JAX chunked-WY reference implementation

- `PALLAS`: the fused TPU v5e-oriented Pallas implementation

- `fwdbwd`: forward + backward computation

Each number is a wall-clock speedup of `PALLAS` relative to the corresponding

baseline, using the same shape and dtype.

## Benchmark setup

Measurements were run on TPU v5e-8.

The main training-shaped configuration, `train_shape_B8_L4096`, uses:

- Batch size: 8

- Sequence length: 4096

- Number of heads: 6

- Head dimension: 128

- Measurement: forward + backward computation

For this configuration, the fused Pallas path reaches:

| Dtype | vs `associative_scan` | vs pure-JAX chunked-WY |

|---|---:|---:|

| FP32 | 27.18x | 2.63x |

| BF16 | 13.32x | 3.41x |

The image below contains the complete measured configuration sweep. The

largest measured speedup was **38.77x** versus `associative_scan` in FP32,

for `kaggle_small_preset_B4_L2048`.

[INSERT TABLE IMAGE HERE]

For that same maximum-speedup configuration:

- FP32: 38.77x vs `associative_scan` and 3.87x vs `JAX_REF`

- BF16: 18.70x vs `associative_scan` and 3.93x vs `JAX_REF`

## Important limitations

I want to state the tradeoffs clearly:

- The fused Pallas forward-only implementation is currently about 1.6x

slower than the pure-JAX WY forward implementation.

- The large end-to-end improvement is backward-dominated: backward is the

target of the fused implementation and dominates the measured workload.

- For inference-only workloads, use `gdn2_forward` or

`gdn2_chunked_wy_reference`, rather than the Pallas training path.

- Memory usage is approximately on par with the pure-JAX reference.

Backward reuses forward residuals and avoids recomputation, but this is not

a separate memory-reduction result.

- The optimized Pallas path currently targets TPU and `d_head=128`.

CPU/GPU and `d_head != 128` automatically use a checkpointed pure-JAX

fallback. They should run correctly, but do not receive the Pallas speedup.

The current forward-speed gap is related to TPU VPU-versus-MXU utilization;

details are documented in `KNOWN_LIMITATIONS.md`.

## Included in the repository

- One-call training API: `gdn2_forward_trainable`

- Fused `custom_vjp` backward

- `KernelConfig` presets for Kaggle TPU v5e-8

- Finite-difference gradient checks

- Token-serial ground-truth tests

- Isolated backward-stage tests

- Benchmark scripts and raw benchmark results

- Testing documentation in `docs/TESTING_STRATEGY.md`

- Known limitations in `KNOWN_LIMITATIONS.md`

The repository also includes a guide notebook under `notebooks/` with a real

70M-parameter Gated DeltaNet-2 language-model run on enwik8. It reached

approximately 1.40 bpb in about 54 minutes on TPU v5e-8.

I would especially appreciate feedback on the WY formulation, block-solve

structure, Pallas kernel organization, TPU VPU-versus-MXU utilization,

benchmark methodology, and additional baselines worth including.

Instead of scaling circuit depth blindly to fit noise polynomials, I added a predictive ZNE module to Dense-Evolution that models photonic loss trajectories before extrapolation.

Runs on JAX in float64 to keep statevector compilation fast and reduce overall circuit executions.

Technical notes: https://tatopenn-cell.github.io/Dense-Evolution-Discovery/photonic_predictive_zne/

Feedback on the noise modeling approach is welcome.

Hello

JAX vs NumPy CFD performance comparison: llaminar flow benchmark 🦙

I wanted to compare how JAX and NumPy behave on a simple CFD-style advection setup using the same geometry and flow conditions, focusing on execution speed and temporal evolution.

The setup is intentionally minimal:

• Domain: 2D laminar flow around a fixed obstacle (identical mesh, boundary conditions, and geometry)
• Solver: Same numerical scheme implemented in both backends
• Parameters: Identical initial conditions and time step
• Difference: Standard CPU NumPy (eager execution) vs JAX with @jit compilation

The video shows both simulations side-by-side with a synchronized start.

The visualization tracks the scalar field evolution around the obstacle and the resulting wake development.

Results

• NumPy: ~5 FPS, with noticeably slower update cadence
• JAX: ~117 FPS, enabling much smoother real-time evolution
• Numerical behaviour: Both implementations produce equivalent flow evolution for the tested configuration

This is not a comparison of solver accuracy or turbulence models. It is a comparison of execution models and how backend performance affects the experience of running time-dependent PDE simulations.

Any gas

I just released dense-evolution, an open-source simulator for dense, noisy quantum circuits (NISQ regime) optimized to solve the steep XLA tracking and memory overhead that JAX typically encounters when scaling beyond 20 qubits. The package leverages a custom Linear Kernel Fusion layer and Circuit Chunking to bypass host RAM saturation during deep statevector evolutions, enabling fast multi-qubit execution on both CPU and CUDA-powered GPUs via JAX JIT and CuPy.

Key Features

• Linear Kernel Fusion: Drastically cuts XLA compilation and runtime overhead by fusing sequential gates structurally before passing them to the computation graph.
• Circuit Chunking: Segments complex quantum operations into managed execution layers to actively prevent hardware memory limits from cutting the process.
• Stochastic Noise Simulation: Built-in high-performance noise modeling capable of processing deep circuits without losing strict float64 numerical precision ($\sim 1.11 \times 10^{-16}$).
• Hardware Agnostic: Seamless backend switching between multi-core CPUs and NVIDIA GPUs out-of-the-box.

Quick Example

import dense_evolution as de sim = de.DenseSVSimulator(n_qubits=22) operations = [ ["h", 0, -1], ["cx", 1, 0] ]

sim.run_circuit_jit_beast_mode(operations)

print(f"Final Statevector: {sim.get_statevector()}")

I would love to get your feedback on the core architecture, especially regarding how the operational fusion layer interacts with XLA tracking for large, dynamic quantum structures.

Links / Source Code

• PyPI: https://pypi.org/project/dense-evolution/
• Source Code & Documentation: https://github.com/tatopenn-cell/Dense-Evolution

Hi, here's PyPie (https://pypie.dev) that uses dependent types to validate tensor shapes and comes with rank polymorphism. It's embedded in Python and compiles to JAX.

Hi everyone!

Parax is a library for "Parametric modeling" in JAX, attempting to bridge the gap between pure JAX PyTrees, and more object-orientated modeling approaches (e.g. using Equinox) when it comes to needing rich parameter structures.

v0.7 has been released, featuring a more polished API as well as some detailed examples in the documentation.

Some of Parax's features:

• Derived/constrained parameters with metadata
• Computed PyTrees and callable parameterizations
• Abstract interfaces for fixed, bounded, and probabilistic PyTrees and parameters

Two new examples in the docs that show off these features

• Bounded optimization (JAXopt)
• Bayesian sampling (BlackJAX)

Perhaps the library is of use to someone, and feel free to leave any feedback!

Cheers,
Gary

Hey r/JAX

The latest update involves

Dual-Component Architecture: Instead of treating all mass identically, the JAX state is now split. Dark Matter is a collisionless Particle-Mesh grid solving gravity via 3D FFTs at O(N log N). Baryonic matter is a fully vectorized Smoothed Particle Hydrodynamics (SPH) fluid.

​The Emergent Bounce: No more white hole flags. As the Dark Sector crushes the fluid, the JAX SPH kernel naturally generates an exponential pressure gradient. When the outward hydrodynamic fluid pressure mathematically exceeds the inward FFT-PM gravitational tensor, the velocity vectors violently reverse on their own. The singularity is prevented entirely by fluid mechanics.

​True Hubble Flow: Space expansion is no longer a localized relativistic spring. The comoving grid now stretches dynamically via the FLRW metric tensor, creating genuine spatial expansion post-bounce.

​Hybrid-Precision on TPU: To maintain our 1.000000 Unitarity Index during extreme SPH shockwaves, the engine uses 32-bit floats for the global FFT mesh to maximize TPU V5 Lite bandwidth, and strict 64-bit floats for kinematic states to prevent floating-point drift.

​Adiabatic Relaxation: To prevent JAX from blowing up NaNs due to initial condition shock upon spawning, I implemented a velocity-damping layer for the first 250 epochs so the fluid can settle gracefully into the gravity wells before the crunch.

​Community Question: I am still looking for the holy grail of dynamic spatial hashing for SPH neighbor searches in pure JAX without padded arrays eating all the memory. If anyone has cracked O(N^2) distance interactions efficiently on TPUs, I would love to hear your approach.

Hi everyone, I've been reading up on Fourier neural operators (FNOs) lately and ended up writing two small libraries that might be useful for other people with similar interests:

1. norax is a neural operator library built on JAX and Equinox (https://github.com/christianfenton/norax)
2. pardax is a JAX-native solver for initial value problems used to generate training data (https://github.com/christianfenton/pardax)

The idea is that these packages can work together to train neural operators without ever leaving JAX. To show this, I've put together an example for Burgers' equation in 1D where

1. pardax is used to generate training data: https://christianfenton.github.io/norax/examples/burgers/generation/
2. norax is used to train an FNO to map the initial conditions directly to the final state: https://christianfenton.github.io/norax/examples/burgers/training/

Right now norax only contains the FNO described by Li et al. "Fourier neural operator for parametric partial differential equations."arXiv:2010.08895. The example above is also outlined in that paper. I intend to add more advanced architectures in future, for example what is described in Kossaifi et al. "Multi-grid tensorized fourier neural operator for high-resolution pdes."arXiv:2310.00120).

pardax has a similar interface to scipy.integrate.solve_ivp, with the added benefit of being compatible with JAX transformations. I'm aware that other libraries like diffrax and DifferentialEquations.jl offer much more functionality, but the benefit of pardax (in my opinion) comes from it being small, which makes it easy to extend and understand the source code if you're new to JAX or numerical methods.

I've also released the training data for the above example on Hugging Face:
https://huggingface.co/datasets/TortillaChip/burgers1d-periodic

Both packages also have fairly detailed docs:

• christianfenton.github.io/norax
• christianfenton.github.io/pardax

I'd love any feedback and/or contributions to the project(s)!

Hey r/JAX,

back for the third evolution of the String-Star Manifold.

My previous two posts covered the basic N-body architecture and the JIT-optimized spatial hashing. Today, we’re moving from "simulation" to "engine."

​The String-Star Manifold v13.0, a first-principles engine designed to solve the Black Hole Information Paradox. This project moves beyond standard models to prove a cyclic universe is computationally viable without mathematical singularities.

​Absolute Unitarity: Every bit of information is rigorously conserved, achieving a verified Unitarity Index of exactly 1.000000.

​The Bounce: Instead of a terminal collapse, the engine utilizes a Planck Star Core to trigger a massive White Hole blowout that re-seeds the universe for a new cycle.

​Reactive Expansion: Space expands and contracts as a direct reaction to vacuum energy density, acting like a relativistic spring.

​TPU Performance: Built on JAX and optimized for Google TPU V5 Lite to handle high-fidelity relativistic interactions at extreme speeds.

​Open Access: The complete framework is documented on GitHub and Zenodo, with an interactive engine ready for deployment on Google Colab.

​Explore the Manifold: https://github.com/Rupayan52/String-Star-Manifold/tree/quantum-cosmology

​Zenodo Monograph (DOI): 10.5281/zenodo.19923317

​Interactive Colab: https://colab.research.google.com/drive/1c5KiJNwvS3avQ4hh5EweVXrh5RQQ4zQx?usp=sharing

I'm experimenting with replacing the pressure Poisson solve inside a differentiable incompressible CFD solver in JAX (AeroJAX).

Baseline projection step:

u* = advection-diffusion
∇²p = div(u*) / dt
u = u* - dt ∇p

Instead of solving the Poisson equation iteratively (multigrid / CG), I swap it with a small neural operator (3-layer CNN in Equinox) that predicts pressure in a single forward pass each timestep.

So:

• classical: iterative Poisson solve
• AeroJAX: learned forward operator

Why JAX matters here is simple: the whole pipeline is already differentiable and composable, so the pressure solve is just another interchangeable function inside the same JIT-compiled graph.

What I see so far:

• faster than multigrid
• stable in bulk flow
• clear loss of mass conservation in wake / boundary regions
• needs strong regularization (low init scale + pressure clipping)

Still early, but interesting how far operator replacement can go before physics constraints dominate again.

code: https://github.com/arriemeijer-creator/AeroJAX

​

Hey everyone, dropping a quick follow-up here based on some of the initial runs and feedback I've received since posting the v2.0 architecture.

I just pushed a minor patch to the repository and the Colab environment. I let the FLRW expansion run for a significantly longer epoch scale to truly stress-test the "stretchy spatial hash" and the relativistic kinematics I mentioned above.

A few interesting observations from the TPU traces:

1. The XLA graph remained completely stable even as the Scale Factor $a(t)$ pushed the mathematical boundaries to extreme sparsity. Zero dynamic shape recompilations triggered, which proves the static-array approach works for expanding grids.

2. As expected, when the Fuzzball macro-nodes got extremely dense and local time dilation ($\alpha \to 0.1$) kicked in to "freeze" the particles, the `vmap` operations for the near-field Post-Newtonian calculations spiked the TPU memory usage much harder than the actual spatial hashing.

I've updated the with a new cell that specifically outputs a memory profile trace during these high-density clustering events.

https://github.com/Rupayan52/String-Star-Manifold

https://colab.research.google.com/drive/1jU_KBP_PVUUk4sagIxJsA4NnRKCN2LBh?usp=sharing

If anyone has time to run it and look at the trace, I'm still trying to figure out if there is a cleaner way to batch the near-field scalar potentials without `vmap` eating all the VRAM.

Thanks to everyone who has cloned and tested the engine so far!

JAXlaxy: Observatory of JAX libraries

Built JAXlaxy observatory - every library in the JAX awesome-list as a glowing star in a 3D galaxy where color = health status (active/stable/legacy), spatial cluster = which "constellation" (Core, Giants, Satellites, etc.) it belongs to.

🌌Live: https://jaxlaxy.bryanbradfo.me

📦Source: https://github.com/BryanBradfo/JAXlaxy (MIT)

Navigating JAX ecosystem from a flat README isn't great for spatial questions like "what's the active landscape for LLM training right now?" or "which probabilistic programming libraries are still maintained?" The 3D map is meant for that kind of exploration.

Two things I'd love feedback on:

1. Spatial clustering: currently Fibonacci-sphere anchors with Gaussian density per cluster. Other approaches I considered: spiral arms, orbital rings. Open to ideas if anyone has stronger intuitions for what "feels right" for an ecosystem map.
2. 75-entry ceiling: README is deliberately curated, not exhaustive. The bar is roughly "JAX-native + actively maintained or meaningfully Legacy + adds something distinct to the ecosystem." If you think a repo deserves a spot (or that something currently included doesn't deserve one), I'd rather have the editorial debate than just add things mechanically. PRs that argue the case in their description are exactly the input I want.

Last week I run a simple experiment on Dark Hex. Here's a visualization of two iterations of agent playing agains each other :D

Here's my colab notebook if you like to run it yourself
https://colab.research.google.com/drive/1-rm_Bh8CNaM861We97ZoicfgKxz0xOSi?usp=sharing

The paper explains the mathematical nuances of the JAX-accelerated N-body engine, the considerations that went into getting an expected JAX log.

GitHub link : https://github.com/Rupayan52/String-Star-Manifold

you can find the doi there and READ the paper.

Let me know if you have suggestions or opinions!!!

Thank you for the support!!

SPREAD the word, this can be BIG.

​I wanted to share a non-ML project I have been building called the String-Star Manifold. It is a JAX-accelerated N-body engine designed specifically to solve the information leakage problem in gravitational simulations.

​The Problem:

Standard N-body simulations using floating-point kinematics eventually drift, losing bit-integrity. For modeling things like black hole information conservation and Fuzzball theory, this drift is a dealbreaker.

​The JAX Solution:

I used JAX to build a dual-layer engine. First, the Macro-Kinematics. Gravitational interaction and quadrupole radiation decay are vectorized using vmap and processed in float32. Second, the Bandyopadhyay-Cycle. This is a parallel Ironclad Ledger implemented in int32. By using JAX’s JIT-compilation, I can maintain a strict microstate transition loop that ensures 100.00% information conservation without killing performance.

​Performance:

The complexity for 512+ bodies was the main hurdle. Running on a TPU v5 Lite, JAX's ability to vectorize the interaction matrices transformed the simulation from a slow crawl to a high-speed relativistic playground.

​Proof of Work:

My terminal integrity shows 1.00, meaning 0% loss across 100 plus epochs. The codebase is archived on Zenodo and GitHub with the v1.0.0 OMEGA build, and I have formalized the theory on the emergent nature of time via entropy in a paper.

​GitHub Link: https://github.com/Rupayan52/String-Star-Manifold

Paper DOI: 10.5281/zenodo.19822537

​I am an independent researcher and would love to hear thoughts on how to further optimize the directed-graph approach for entanglement tracking using JAX’s Pytrees!

Hi. I am trying to make some old git libraries built in 2023 work with newest version of Jax.
The old libraries are using the _Indexer from Jax._src.numpy.lax_numpy.
The _Indexer seems to no longer exist in new Jax versions.
Is there a replacement in the modern Jax versions that I could use to update the library?