RichardGrumitt
Follow
Stars
All sky Synchrotron emission simulation with JAX
Curated list of astronomy/astrophysics code packages using JAX
JAX Tutorial notebooks : basics, crash & tips, usage of optax/JaxOptim/Numpyro
Sliced Iterative Normalizing Flow with Minimal Dependencies
Differentiable and accelerated wavelet transform on the sphere with JAX
Interpolation and function approximation with JAX
A Gaussian Process based 21-cm foreground separation code.
List the AI for Science papers accepted by top conferences
State of the art inference for your bayesian models.
Implements Optimization and approximate uncertainty quantification algorithms, Ensemble Kalman Inversion, and Ensemble Kalman Processes.
XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML
Modern normalizing flows in Python. Simple to use and easily extensible.
A differentiable PDE solving framework for machine learning
All the handwritten notes 📝 and source code files 🖥️ used in my YouTube Videos on Machine Learning & Simulation (https://www.youtube.com/channel/UCh0P7KwJhuQ4vrzc3IRuw4Q)
Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.
Optimal transport tools implemented with the JAX framework, to get differentiable, parallel and jit-able computations.
Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/
Additive manufacturing simulation with JAX.
Sliced Iterative Generator (SIG) & Gaussianizing Iterative Slicing (GIS)
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
Normalizing-flow enhanced sampling package for probabilistic inference in Jax
MCHMC: sampler from an arbitrary differentiable distribution
Second Order Optimization and Curvature Estimation with K-FAC in JAX.
Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python ⚡
JAX - A curated list of resources https://github.com/google/jax