JAXsaw: Visualizing Composable Neural Architectures in JAX

JAXsaw is a lightweight, GPU-accelerated framework for building and exploring neural networks in JAX. Its modular building blocks support architectures ranging from dense networks and CNNs to transformers, autoencoders, and GANs.

The library supports model development, training, and evaluation, with documented examples and visualizations for learning and experimentation.

• Installation
• User guide
  • Core concepts
  • State, checkpoints, and reproducibility
  • Examples and datasets
  • Compatibility and limitations
• Architecture tutorials
  • MLP: Titanic Classification
  • CNN: MNIST Classification
  • RNN: Seasonal Temperature Forecasting
  • LSTM: Market Volatility Forecasting
  • Transformer: Language Modeling
  • Autoencoder: MNIST Reconstruction
  • GAN: Fashion-MNIST Generation
• API tutorials
  • Modules and parameters
  • Training and evaluation
  • Data and batching
  • Checkpoints and randomness
  • Inspection and custom layers
• API reference
  • API and semantics
  • Python reference