Equinox: neural networks in JAX via callable PyTrees and filtered transformations

30 October 2021

Papers citing "Equinox: neural networks in JAX via callable PyTrees and filtered transformations"

50 / 56 papers shown

Title
EconoJax: A Fast & Scalable Economic Simulation in Jax Koen Ponse Aske Plaat N. V. Stein Thomas M. Moerland 28 0 0 20 May 2025
Binding threshold units with artificial oscillatory neurons V. Fanaskov Ivan V. Oseledets 49 0 0 06 May 2025
Energy-Based Coarse-Graining in Molecular Dynamics: A Flow-Based Framework Without Data Maximilian Stupp P. S. Koutsourelakis 40 0 0 29 Apr 2025
Uncertainty quantification of neural network models of evolving processes via Langevin sampling C. Safta Reese E. Jones Ravi G. Patel Raelynn Wonnacot Dan S. Bolintineanu Craig M. Hamel S. Kramer BDL 36 0 0 21 Apr 2025
Scaling sparse feature circuit finding for in-context learning Dmitrii Kharlapenko S. Kamath S Fazl Barez Arthur Conmy Neel Nanda 26 0 0 18 Apr 2025
Modelling Mean-Field Games with Neural Ordinary Differential Equations Anna C. M. Thöni Yoram Bachrach Tal Kachman 38 0 0 17 Apr 2025
Scaling Up Resonate-and-Fire Networks for Fast Deep Learning Thomas E. Huber J. Lecomte Borislav Polovnikov Axel von Arnim 36 0 0 01 Apr 2025
Neural ODE Transformers: Analyzing Internal Dynamics and Adaptive Fine-tuning Anh Tong Thanh Nguyen-Tang Dongeun Lee Duc Nguyen Toan M. Tran David Hall Cheongwoong Kang Jaesik Choi 35 0 0 03 Mar 2025
PRDP: Progressively Refined Differentiable Physics Kanishk Bhatia Felix Koehler Nils Thuerey AI4CE 63 0 0 26 Feb 2025
Per-channel autoregressive linear prediction padding in tiled CNN processing of 2D spatial data Olli Niemitalo Otto Rosenberg Nathaniel Narra Olli Koskela Iivari Kunttu 69 0 0 17 Feb 2025
Benchmarking Predictive Coding Networks -- Made Simple Luca Pinchetti Chang Qi Oleh Lokshyn Gaspard Olivers Cornelius Emde ... Simon Frieder Bayar I. Menzat Rafal Bogacz Thomas Lukasiewicz Tommaso Salvatori 127 5 0 17 Feb 2025
Modelling Chemical Reaction Networks using Neural Ordinary Differential Equations Anna C. M. Thöni William Robinson Yoram Bachrach Wilhelm Huck Tal Kachman 48 0 0 11 Feb 2025
Biomechanical Reconstruction with Confidence Intervals from Multiview Markerless Motion Capture R. J. Cotton Fabian H. Sinz 49 0 0 10 Feb 2025
Towards Foundational Models for Dynamical System Reconstruction: Hierarchical Meta-Learning via Mixture of Experts Roussel Desmond Nzoyem David A.W. Barton Tom Deakin 72 1 0 07 Feb 2025
Optimizing Automatic Differentiation with Deep Reinforcement Learning Jamie Lohoff Emre Neftci 58 1 0 28 Jan 2025
Learning state and proposal dynamics in state-space models using differentiable particle filters and neural networks Benjamin Cox Santiago Segarra Victor Elvira 73 0 0 23 Nov 2024
Adjoint-based online learning of two-layer quasi-geostrophic baroclinic turbulence Fei Er Yan Hugo Frezat Julien Le Sommer Julian Mak Karl Otness 62 1 0 21 Nov 2024
Constrained Latent Action Policies for Model-Based Offline Reinforcement Learning Marvin Alles Philip Becker-Ehmck Patrick van der Smagt Maximilian Karl OffRL 34 0 0 07 Nov 2024
Physics-informed neural networks viewpoint for solving the Dyson-Schwinger equations of quantum electrodynamics Rodrigo Carmo Terin PINN 64 1 0 04 Nov 2024
Adaptive Data Optimization: Dynamic Sample Selection with Scaling Laws Yiding Jiang Allan Zhou Zhili Feng Sadhika Malladi J. Zico Kolter 39 15 0 15 Oct 2024
pyhgf: A neural network library for predictive coding Nicolas Legrand Lilian A. E. Weber Peter Thestrup Waade Anna Hedvig Møller Daugaard Mojtaba Khodadadi Nace Mikuš C. Mathys AI4CE 23 0 0 11 Oct 2024
DeepOSets: Non-Autoregressive In-Context Learning of Supervised Learning Operators Shao-Ting Chiu Junyuan Hong Ulisses Braga-Neto BDL 28 0 0 11 Oct 2024
Learning responsibility allocations for multi-agent interactions: A differentiable optimization approach with control barrier functions Isaac Remy David Fridovich-Keil Karen Leung 19 0 0 09 Oct 2024
Automatic debiasing of neural networks via moment-constrained learning Christian L. Hines Oliver J. Hines CML OOD 31 0 0 29 Sep 2024
Let the Quantum Creep In: Designing Quantum Neural Network Models by Gradually Swapping Out Classical Components Peiyong Wang Casey. R. Myers Lloyd C. L. Hollenberg U. Parampalli 137 0 0 26 Sep 2024
Neural Differential Appearance Equations Chen Liu Tobias Ritschel 31 0 0 23 Sep 2024
Inverse decision-making using neural amortized Bayesian actors Dominik Straub Tobias F. Niehues Jan Peters Constantin Rothkopf 28 0 0 04 Sep 2024
Certified Robust Invariant Polytope Training in Neural Controlled ODEs Akash Harapanahalli Samuel Coogan 42 2 0 02 Aug 2024
Neural Octahedral Field: Octahedral prior for simultaneous smoothing and sharp edge regularization Gareth Seneque Tao Yu 21 0 0 01 Aug 2024
Penzai + Treescope: A Toolkit for Interpreting, Visualizing, and Editing Models As Data Mingshu Li 36 3 0 01 Aug 2024
NeuroSEM: A hybrid framework for simulating multiphysics problems by coupling PINNs and spectral elements K. Shukla Zongren Zou Chi Hin Chan Additi Pandey Zhicheng Wang George Karniadakis PINN 52 7 0 30 Jul 2024
SoftCVI: Contrastive variational inference with self-generated soft labels Daniel Ward Mark Beaumont Matteo Fasiolo BDL 53 0 0 22 Jul 2024
Separable Operator Networks Xinling Yu S. Hooten Z. Liu Yequan Zhao M. Fiorentino T. Van Vaerenbergh Zheng-Wei Zhang 51 4 0 15 Jul 2024
Engineering morphogenesis of cell clusters with differentiable programming Ramya Deshpande Francesco Mottes Ariana-Dalia Vlad Michael P. Brenner Alma dal Co AI4CE 29 0 0 08 Jul 2024
Learning deformable linear object dynamics from a single trajectory Shamil Mamedov A. R. Geist Ruan Viljoen Sebastian Trimpe Jan Swevers PINN AI4CE 32 1 0 03 Jul 2024
Meta-Learning an Evolvable Developmental Encoding Milton L. Montero Erwan Plantec Eleni Nisioti J. Pedersen Sebastian Risi 32 0 0 13 Jun 2024
ConDiff: A Challenging Dataset for Neural Solvers of Partial Differential Equations Vladislav Trifonov Alexander Rudikov Oleg Iliev Ivan V. Oseledets Ekaterina A. Muravleva Ekaterina Muravleva DiffM AI4CE 33 0 0 07 Jun 2024
Astral: training physics-informed neural networks with error majorants V. Fanaskov Tianchi Yu Alexander Rudikov Ivan V. Oseledets 33 1 0 04 Jun 2024
Fusing uncalibrated IMUs and handheld smartphone video to reconstruct knee kinematics J. D. Peiffer Kunal Shah Shawana Anarwala Kayan Abdou R. J. Cotton 43 2 0 27 May 2024
Learning from Linear Algebra: A Graph Neural Network Approach to Preconditioner Design for Conjugate Gradient Solvers Vladislav Trifonov Alexander Rudikov Oleg Iliev Ivan V. Oseledets Ekaterina A. Muravleva Ekaterina Muravleva 39 2 0 24 May 2024
Scalable physical source-to-field inference with hypernetworks Berian James Stefan Pollok Ignacio Peis J. Frellsen Rasmus Bjørk AI4CE 38 0 0 07 May 2024
Tripod: Three Complementary Inductive Biases for Disentangled Representation Learning Kyle Hsu Jubayer Ibn Hamid Kaylee Burns Chelsea Finn Jiajun Wu CML 26 4 0 16 Apr 2024
Differentiable Biomechanics Unlocks Opportunities for Markerless Motion Capture R. J. Cotton 34 5 0 27 Feb 2024
Re-Envisioning Numerical Information Field Theory (NIFTy.re): A Library for Gaussian Processes and Variational Inference G. Edenhofer Philipp Frank Jakob Roth R. Leike Massin Guerdi L. Scheel-Platz M. Guardiani Vincent Eberle M. Westerkamp T. Ensslin 28 9 0 26 Feb 2024
Scaling physics-informed hard constraints with mixture-of-experts N. Chalapathi Yiheng Du Aditi Krishnapriyan AI4CE 32 12 0 20 Feb 2024
Align Your Intents: Offline Imitation Learning via Optimal Transport Maksim Bobrin N. Buzun Dmitrii Krylov Dmitry V. Dylov OffRL 43 3 0 20 Feb 2024
Glocal Hypergradient Estimation with Koopman Operator Ryuichiro Hataya Yoshinobu Kawahara 39 2 0 05 Feb 2024
Enhancing Score-Based Sampling Methods with Ensembles T. Bischoff B. Riel 17 1 0 31 Jan 2024
$$\texttt{immrax}$: A Parallelizable and Differentiable Toolbox for Interval Analysis and Mixed Monotone Reachability in JAX$ $\texttt{immrax}$ : A Parallelizable and Differentiable Toolbox for Interval Analysis and Mixed Monotone Reachability in JAX Akash Harapanahalli Saber Jafarpour Samuel Coogan 24 3 0 21 Jan 2024
XLB: A differentiable massively parallel lattice Boltzmann library in Python Mohammadmehdi Ataei H. Salehipour AI4CE 22 13 0 27 Nov 2023