JAX

exojaxHajimeKawahara/exojax

A JAX-based differentiable spectral modeling library for exoplanets, brown dwarfs, and M dwarfs.

KernexASEM000/kernex

Kernex extends JAX with kmap and kscan for differentiable stencil computations, enabling efficient array transformations.

GNNs for Solving Combinatorial Optimization ProblemsIvanIsCoding/GNN-for-Combinatorial-Optimization

A JAX + Flax implementation of physics-inspired graph neural networks for solving combinatorial optimization problems like Max-Cut and Maximum Independent Set.

delta PVromanodev/deltapv

A photovoltaic simulator with automatic differentiation for solar cell modeling and optimization, built on JAX.

PGMaxvicariousinc/PGMax

A JAX-based library for loopy belief propagation on discrete factor graphs, enabling efficient probabilistic inference.

astronomixleo1200/astronomix

A differentiable hydrodynamics and magnetohydrodynamics code for astrophysics built with JAX, enabling gradient-based inverse modeling and multi-GPU simulations.

JAXFitdipolar-quantum-gases/jaxfit

GPU/TPU accelerated nonlinear least-squares curve fitting using JAX, designed as a drop-in replacement for SciPy's curve_fit.

DiffeRTjeertmans/DiffeRT

A differentiable ray tracing toolbox for radio propagation simulations, built on JAX for optimization and machine learning.