Welcome to blox’s documentation!

blox is a functional and lightweight neural network library for JAX.

The entire mental model fits in one line:

outputs, params = model(params, inputs)

Parameters go in, outputs and updated parameters come out. Because state flows explicitly through your code, all JAX transformations work out of the box. No wrappers, no decorators, no surprises.

Getting Started

• Installation
  • Prerequisites
  • Install from PyPI
  • Install from Source
• MNIST Classification with blox
• Recurrent Neural Networks: From Sine Waves to Shakespeare
  • Part 1: Modeling a Sine Wave (Regression)
  • Part 2: Tiny Shakespeare (Generation)

Key Concepts

• Design Principles
  • The Mental Model
  • Functional Purity
  • Core Abstractions
  • The Params Container
  • JAX Compatibility
  • Training Pattern
  • Lazy Initialization
  • Why the Verbosity?
  • Decoupled Params and Graph

Advanced Topics

• Checkpointable Training on Google Colab
  • The Problem
  • The Solution
  • The Functional Approach
  • How to Use This Notebook
• 1. Setup
• 2. Configuration
• 3. Utilities
• 4. Resumable Dataset
• 5. Model
• 6. Training Step
• 7. Initialize and Restore
• 8. Training Loop
• 9. Evaluation
• 10. Summary
  • What We Checkpoint
  • Key Patterns
  • References
• Two Approaches
  • Approach 1: Monkey-patching an Existing Model
• 1. Define an MLP Module
• 2. Explore the Graph
• 3. The LoRA Pattern
• 4. Apply LoRA to Selected Layers
• 5. Initialize LoRA Parameters
• 6. Freeze Base Weights
• 7. Training with LoRA
• 8. Verify Only LoRA Weights Changed
• 9. Merging LoRA Weights (Optional)
  • Approach 2: LoRA-aware Model Design
• Create a LoRA-aware MLP
• Benefits of the LoRA-aware Approach
• Toggling LoRA
• Sharp Bits in blox
  • Contents
• 1. The Params Container
  • Locked vs Unlocked
  • Immutability
  • Split and Merge for Gradient Computation
  • Params Ownership and Memory
• 2. RNG Handling in Parallel Contexts
  • How blox Handles Randomness
  • The Problem: Replicated RNG State
  • The Solution: Fold in the Lane Index
  • Why Restore the Original Seed?
• 3. Init vs Runtime
  • The is_locked Pattern
  • Sharded Initialization: JIT vs shard_map
• 4. Graph and Module Modifications
  • The Problem: JIT Caching and Module State
  • blox Philosophy: Prefer Static Models
  • Pattern 1: Build Separate Static Models
  • Pattern 2: Create Separate Functions
  • Sharing Params Across Models
• 5. Summary
  • Sharp Edges and Patterns
  • Key Design Principles
  • References

Reference

• API Reference
  • Core Interfaces
  • Layers
  • Pooling
  • Sequence Processing
  • Visualization

Indices and tables

• Index

• Module Index

• Search Page