English

SymJAX: symbolic CPU/GPU/TPU programming

Mathematical Software 2020-05-22 v1 Computer Vision and Pattern Recognition
Authors: Randall Balestriero
View on arXiv ↗ PDF ↗

Abstract

SymJAX is a symbolic programming version of JAX simplifying graph input/output/updates and providing additional functionalities for general machine learning and deep learning applications. From an user perspective SymJAX provides a la Theano experience with fast graph optimization/compilation and broad hardware support, along with Lasagne-like deep learning functionalities.

Keywords

symbolic execution compiler optimization

Cite

@article{arxiv.2005.10635,
  title  = {SymJAX: symbolic CPU/GPU/TPU programming},
  author = {Randall Balestriero},
  journal= {arXiv preprint arXiv:2005.10635},
  year   = {2020}
}

Related papers

View all related →
Computational Physics · Physics
msmJAX: Fast and Differentiable Electrostatics on the GPU in Python
Florian Buchner, Johannes Schörghuber, Nico Unglert, Jesús Carrete +1
2025-10-08
Quantum Physics · Physics
SymPhase: Phase Symbolization for Fast Simulation of Stabilizer Circuits
Wang Fang, Mingsheng Ying
2023-11-23
Mathematical Software · Computer Science
BlackJAX: Composable Bayesian inference in JAX
Alberto Cabezas, Adrien Corenflos, Junpeng Lao, Rémi Louf +18
2024-02-23
Neural and Evolutionary Computing · Computer Science
SNNAX -- Spiking Neural Networks in JAX
Jamie Lohoff, Jan Finkbeiner, Emre Neftci
2025-01-28
Artificial Intelligence · Computer Science
PuzzleJAX: A Benchmark for Reasoning and Learning
Sam Earle, Graham Todd, Yuchen Li, Ahmed Khalifa +4
2025-08-26
Neural and Evolutionary Computing · Computer Science
EvoJAX: Hardware-Accelerated Neuroevolution
Yujin Tang, Yingtao Tian, David Ha
2022-04-07
Computation and Language · Computer Science
High-performance symbolic-numerics via multiple dispatch
Shashi Gowda, Yingbo Ma, Alessandro Cheli, Maja Gwozdz +3
2022-02-08
Computational Physics · Physics
SymPhas --General purpose software for phase-field, phase-field crystal and reaction-diffusion simulations
Steven A. Silber, Mikko Karttunen
2021-09-07
Robotics · Computer Science
SymForce: Symbolic Computation and Code Generation for Robotics
Hayk Martiros, Aaron Miller, Nathan Bucki, Bradley Solliday +12
2022-11-29
Computational Physics · Physics
Parallel and GPU accelerated code for phase-field and reaction-diffusion simulations
Steven A. Silber, Mikko Karttunen
2025-11-14
Computational Physics · Physics
TMMax: High-performance modeling of multilayer thin-film structures using transfer matrix method with JAX
Bahrem Serhat Danis, Esra Zayim
2025-10-27
Distributed, Parallel, and Cluster Computing · Computer Science
Synkhronos: a Multi-GPU Theano Extension for Data Parallelism
Adam Stooke, Pieter Abbeel
2017-10-13
Instrumentation and Methods for Astrophysics · Physics
$\texttt{synax}$: A Differentiable and GPU-accelerated Synchrotron Simulation Package
Kangning Diao, Zack Li, Richard D. P. Grumitt, Yi Mao
2025-03-26
Distributed, Parallel, and Cluster Computing · Computer Science
DrJAX: Scalable and Differentiable MapReduce Primitives in JAX
Keith Rush, Zachary Charles, Zachary Garrett, Sean Augenstein +1
2024-07-19
High Energy Physics - Phenomenology · Physics
Differentiable Matrix Elements with MadJax
Lukas Heinrich, Michael Kagan
2023-03-01
Logic in Computer Science · Computer Science
Denotational Semantics and a Fast Interpreter for jq
Michael Färber
2023-02-22
Logic in Computer Science · Computer Science
SCJ-Circus: a refinement-oriented formal notation for Safety-Critical Java
Alvaro Miyazawa, Ana Cavalcanti
2016-06-08
Programming Languages · Computer Science
Interlanguages and synchronic models of computation
Alexander Victor Berka
2010-05-31
Distributed, Parallel, and Cluster Computing · Computer Science
SIMD-X: Programming and Processing of Graph Algorithms on GPUs
Hang Liu, H. Howie Huang
2018-12-12
Combinatorics · Mathematics
simpcomp -- A GAP toolbox for simplicial complexes
Felix Effenberger, Jonathan Spreer
2019-10-24
Logic in Computer Science · Computer Science
Towards Specifying Symbolic Computation
Jacques Carette, William M. Farmer
2019-05-07
Chemical Physics · Physics
Apax: A Flexible and Performant Framework For The Development of Machine-Learned Interatomic Potentials
Moritz René Schäfer, Nico Segreto, Fabian Zills, Christian Holm +1
2025-11-19
Programming Languages · Computer Science
Prism: Symbolic Superoptimization of Tensor Programs
Mengdi Wu, Xiaoyu Jiang, Oded Padon, Zhihao Jia
2026-04-17
Software Engineering · Computer Science
SysMap: A Lightweight Software Visualization Tool to Analyze the Software Evolution of a System
Fazle Rabbi, Nishat Tasnim Niloy, Nadia Nahar, Md. Nurul Ahad Tawhid
2021-08-16
Machine Learning · Computer Science
SymFlux: deep symbolic regression of Hamiltonian vector fields
M. A. Evangelista-Alvarado, P. Suárez-Serrato
2025-07-10
R2 v1 2026-06-23T15:42:57.238Z