How machine-moon/tinax works
Tinax is a typed library of explicit JAX-ecosystem productivity primitives, not a single application. The user selects and calls the module for a task; key outcomes include transformed arrays, deterministic batches, compiled or differentiated functions, distributed arrays, checkpoints, and interoperable weights. The graph groups its major capability areas and shows internal calls where the excerpts or SOURCE INDEX establish them. Many modules are represented from README capabilities and the file tree; their detailed wiring was not sampled.
Main entry points: _safetensors.py Safetensors interchange · _training.py Training checkpoints · _models.py Weight models · _validation.py Weight validation
Components 26
JAX primitives
- Array conversion and inspection
tinax/array/_conversion.py - Device discovery and validation
tinax/device/_detect.py - Typed RNG keys
tinax/random/_keys.py - Autodiff transforms
tinax/grad/_transforms.py - Autodiff validation
tinax/grad/_validation.py - Compiled batching
tinax/jit/_batching.py - Trace-budgeted compilation
tinax/jit/_compile.py - Compilation validation
tinax/jit/_validation.py - Debug observation
tinax/debug - NNX graph copies
tinax/nn/_graph.py
Data and checkpoints
- Deterministic input pipelines
tinax/data/_pipelines.py - Batch transforms
tinax/data/_generic.py - Training checkpoints
tinax/checkpointing/_training.py - Checkpoint operations
tinax/checkpointing/_operations.py - Checkpoint naming
tinax/checkpointing/_names.py
Parallel computing
- Distributed arrays
tinax/parallel/_arrays.py - Device meshes
tinax/parallel/_mesh.py - Manual SPMD mapping
tinax/parallel/_shard_map.py - Layout validation
tinax/parallel/_validation.py
Weight interchange
- Safetensors interchange
tinax/weights/_safetensors.py - Weight file handling
tinax/weights/_files.py - Dtype and format rules
tinax/weights/_formats.py - Weight models
tinax/weights/_models.py - Weight validation
tinax/weights/_validation.py
Application boundaries
- CLI and logging boundaries
tinax/stdlib
Other components
- Library user
Connections 32
- Library user → Array conversion and inspection: converts arrays
README.md - Library user → Device discovery and validation: inspects devices
README.md - Library user → Typed RNG keys: derives keys
README.md - Library user → Autodiff transforms: differentiates functions
README.md - Library user → Compiled batching: batches functions
README.md - Library user → Trace-budgeted compilation: compiles functions
README.md - Library user → Debug observation: observes execution
README.md - Library user → NNX graph copies: copies graphs
README.md - Library user → Deterministic input pipelines: builds input pipelines
README.md - Library user → Training checkpoints: saves training state
README.md - Library user → Distributed arrays: constructs distributed arrays
README.md - Library user → Device meshes: configures parallelism
README.md - Library user → Manual SPMD mapping: maps SPMD work
README.md - Library user → Safetensors interchange: interchanges weights
README.md - Library user → CLI and logging boundaries: sets app boundaries
README.md - Autodiff transforms → Autodiff validation: validates arguments
tinax/grad/_transforms.py - Compiled batching → Compilation validation: validates settings
tinax/jit/_batching.py - Trace-budgeted compilation → Compilation validation: validates settings
tinax/jit/_compile.py - Deterministic input pipelines → Batch transforms: uses batch transform
tinax/data/_pipelines.py - Training checkpoints → Checkpoint operations: loads and saves
tinax/checkpointing/_training.py - Training checkpoints → Checkpoint naming: resolves names
tinax/checkpointing/_training.py - Checkpoint operations → Checkpoint naming: validates names
tinax/checkpointing/_operations.py - Distributed arrays → Layout validation: validates layouts
tinax/parallel/_arrays.py - Device meshes → Layout validation: validates mesh
tinax/parallel/_mesh.py - Manual SPMD mapping → Layout validation: validates layout
- Safetensors interchange → Weight file handling: handles files
tinax/weights/_safetensors.py - Safetensors interchange → Dtype and format rules: uses dtype rules
tinax/weights/_safetensors.py - Safetensors interchange → Weight models: uses data models
tinax/weights/_safetensors.py - Safetensors interchange → Weight validation: validates inputs
tinax/weights/_safetensors.py - Weight models → Dtype and format rules: uses dtype rules
tinax/weights/_models.py - Weight models → Weight validation: validates values
tinax/weights/_models.py - Weight validation → Dtype and format rules: checks serialization
tinax/weights/_validation.py