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"

29 / 29 papers shown

Title
Training a Foundation Model for Materials on a Budget Teddy Koker Tess Smidt 0 0 0 22 Aug 2025
NiceWebRL: a Python library for human subject experiments with reinforcement learning environments Wilka Carvalho Vikram Goddla Ishaan Sinha Hoon Shin Kunal Jha OffRL 20 0 0 21 Aug 2025
Weight-Space Linear Recurrent Neural Networks Roussel Desmond Nzoyem Nawid Keshtmand Idriss Tsayem David A.W. Barton Tom Deakin 113 0 0 01 Jun 2025
FlowQ: Energy-Guided Flow Policies for Offline Reinforcement Learning Marvin Alles Nutan Chen Patrick van der Smagt Botond Cseke 155 0 0 20 May 2025
Uncertainty quantification of neural network models of evolving processes via Langevin sampling Cosmin Safta Reese E. Jones Ravi G. Patel Raelynn Wonnacot Dan S. Bolintineanu Craig M. Hamel S. Kramer BDL 185 1 0 21 Apr 2025
Modelling Mean-Field Games with Neural Ordinary Differential Equations Anna C. M. Thöni Yoram Bachrach Tal Kachman 141 0 0 17 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 192 4 0 03 Mar 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 121 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 246 13 0 17 Feb 2025
Neural Flow Samplers with Shortcut Models Wuhao Chen Zijing Ou Yingzhen Li 324 2 0 11 Feb 2025
Biomechanical Reconstruction with Confidence Intervals from Multiview Markerless Motion Capture R. J. Cotton Fabian H. Sinz 172 2 0 10 Feb 2025
Towards Foundational Models for Dynamical System Reconstruction: Hierarchical Meta-Learning via Mixture of Experts Roussel Desmond Nzoyem Grant Stevens Amarpal Sahota David A.W. Barton Tom Deakin 243 4 0 07 Feb 2025
Optimizing Automatic Differentiation with Deep Reinforcement Learning Jamie Lohoff Emre Neftci 227 3 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 231 0 0 23 Nov 2024
Constrained Latent Action Policies for Model-Based Offline Reinforcement Learning Marvin Alles Philip Becker-Ehmck Patrick van der Smagt Maximilian Karl OffRL 203 2 0 07 Nov 2024
Physics-informed neural networks viewpoint for solving the Dyson-Schwinger equations of quantum electrodynamics Rodrigo Carmo Terin PINN 160 2 0 04 Nov 2024
EconoJax: A Fast & Scalable Economic Simulation in Jax Koen Ponse Aske Plaat Niki van Stein Thomas M. Moerland 126 1 0 29 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 76 0 0 11 Oct 2024
DeepOSets: Non-Autoregressive In-Context Learning of Supervised Learning Operators Shao-Ting Chiu Junyuan Hong Ulisses Braga-Neto BDL 87 1 0 11 Oct 2024
Automatic debiasing of neural networks via moment-constrained learning Christian L. Hines Oliver J. Hines CML OOD 218 1 0 29 Sep 2024
Neural Differential Appearance Equations Chen Liu Tobias Ritschel 113 1 0 23 Sep 2024
Inverse decision-making using neural amortized Bayesian actors Dominik Straub Tobias F. Niehues Jan Peters Constantin Rothkopf 222 1 0 04 Sep 2024
SoftCVI: Contrastive variational inference with self-generated soft labels Daniel Ward Mark Beaumont Matteo Fasiolo BDL 321 1 0 22 Jul 2024
Engineering morphogenesis of cell clusters with differentiable programming Ramya Deshpande Francesco Mottes Ariana-Dalia Vlad Michael P. Brenner Alma dal Co AI4CE 102 1 0 08 Jul 2024
ConDiff: A Challenging Dataset for Neural Solvers of Partial Differential Equations Vladislav Trifonov Alexander Rudikov Oleg Iliev Ivan Oseledets Ekaterina Muravleva Ekaterina Muravleva DiffM AI4CE 106 0 0 07 Jun 2024
Learning from Linear Algebra: A Graph Neural Network Approach to Preconditioner Design for Conjugate Gradient Solvers Vladislav Trifonov Alexander Rudikov Oleg Iliev Ivan Oseledets Ekaterina Muravleva Ekaterina Muravleva 128 6 0 24 May 2024
Gradient-Free Score-Based Sampling Methods with Ensembles T. Bischoff B. Riel 92 1 0 31 Jan 2024
Differentiable DG with Neural Operator Source Term Correction Shinhoo Kang Emil M. Constantinescu AI4CE 142 0 0 29 Oct 2023
Robust Neural Posterior Estimation and Statistical Model Criticism Daniel Ward Patrick W Cannon Mark Beaumont Matteo Fasiolo Sebastian M. Schmon 106 44 0 12 Oct 2022