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

Kidger, Patrick; Garcia, Cristian

doi:10.48550/arxiv.2111.00254

Cited by 3 publications

(2 citation statements)

References 5 publications

(6 reference statements)

Supporting

Mentioning

Contrasting

Order By: Relevance

“…( 2) Operator overloading, which computes derivatives in real time using operators coded in classes. Some Python libraries that implement this technique are Auto_diff (Nobel, 2020), Autograd (Maclaurin et al, 2015), JAX (Kidger and Garcia, 2021), Chainer (Tokui et al, 2019) and PyTorch (Paszke et al, 2017).…”

Section: Methodsmentioning

confidence: 99%

Optimization on frequency constraints with FFT using automatic differentiation on hybrid ODE applications

Agobert,

Delinchant,

Gerbaud

2024

COMPEL

View full text Add to dashboard Cite

Purpose This study aims to optimize electrical systems represented by ordinary differential equations and events, using their frequency spectrum is an important purpose for designers, especially to calculate harmonics. Design/methodology/approach This paper presents a methodology to achieve this, by using a gradient-based optimization algorithm. The paper proposes to use a time simulation of the electrical system, and then to compute its frequency spectrum in the optimization loop. Findings The paper shows how to proceed efficiently to compute the frequency spectrum of an electrical system to include it in an optimization loop. Derivatives of the frequency spectrum such as the optimization inputs can also be calculated. This is possible even if the sized system behavior cannot be defined a priori, e.g. when there are static converters or electrical devices with natural switching. Originality/value Using an efficient sequential quadratic programming optimizer, automatic differentiation is used to compute the model gradients. Frequency spectrum derivatives with respect to the optimization inputs are calculated by an analytical formula. The methodology uses a “white-box” approach so that automatic differentiation and the differential equations simulator can be used, unlike most state-of-the-art simulators.

show abstract

Section: Methodsmentioning

confidence: 99%

Optimization on frequency constraints with FFT using automatic differentiation on hybrid ODE applications

Agobert,

Delinchant,

Gerbaud

2024

COMPEL

View full text Add to dashboard Cite

show abstract

“…The code is written in Python and utilizes JAX [6] for performing automatic differentiation and numerical optimization for model fitting. All models were implemented from scratch, and we utilize equinox [17] for maintaining readability and elegance of the code. In our code base repository, we provide the code and scripts for reproducing all results in this paper.…”

Section: Supplementary Materialsmentioning

confidence: 99%

Bayesian nonparametric (non-)renewal processes for analyzing neural spike train variability

Liu,

Lengyel

2023

Preprint

View full text Add to dashboard Cite

Neural spiking activity is generally variable, non-stationary, and exhibits complex dependencies on covariates, such as sensory input or behavior. These dependencies have been proposed to be signatures of specific computations, and so characterizing them with quantitative rigor is critical for understanding neural computations. Approaches based on point processes provide a principled statistical framework for modeling neural spiking activity. However, currently, they only allow the instantaneous mean, but not the instantaneous variability, of responses to depend on covariates. To resolve this limitation, we propose a scalable Bayesian approach generalizing modulated renewal processes using sparse variational Gaussian processes. We leverage pathwise conditioning for computing nonparametric priors over conditional interspike interval distributions and rely on automatic relevance determination to detect lagging interspike interval dependencies beyond renewal order. After systematically validating our method on synthetic data, we apply it to two foundational datasets of animal navigation: head direction cells in freely moving mice and hippocampal place cells in rats running along a linear track. Our model exhibits competitive or better predictive power compared to state-of-the-art baselines, and outperforms them in terms of capturing interspike interval statistics. These results confirm the importance of modeling covariate-dependent spiking variability, and further analyses of our fitted models reveal rich patterns of variability modulation beyond the temporal resolution of flexible count-based approaches.

show abstract

Neural Sequences Underlying Directed Turning inC. elegans

Kramer,

Wan,

Pugliese

et al. 2024

Preprint

View full text Add to dashboard Cite

Complex behaviors like navigation rely on sequenced motor outputs that combine to generate effective movement. The brain-wide organization of the circuits that integrate sensory signals to select and execute appropriate motor sequences is not well understood. Here, we characterize the architecture of neural circuits that controlC. elegansolfactory navigation. We identify error-correcting turns during navigation and use whole-brain calcium imaging and cell-specific perturbations to determine their neural underpinnings. These turns occur as motor sequences accompanied by neural sequences, in which defined neurons activate in a stereotyped order during each turn. Distinct neurons in this sequence respond to sensory cues, anticipate upcoming turn directions, and drive movement, linking key features of this sensorimotor behavior across time. The neuromodulator tyramine coordinates these sequential brain dynamics. Our results illustrate how neuromodulation can act on a defined neural architecture to generate sequential patterns of activity that link sensory cues to motor actions.

show abstract

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

Cited by 3 publications

References 5 publications

Optimization on frequency constraints with FFT using automatic differentiation on hybrid ODE applications

Optimization on frequency constraints with FFT using automatic differentiation on hybrid ODE applications

Bayesian nonparametric (non-)renewal processes for analyzing neural spike train variability

Neural Sequences Underlying Directed Turning inC. elegans

Contact Info

Product

Resources

About