# machine-moon/tinax architecture

> Architecture diagram of the GitHub repository machine-moon/tinax, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-29). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/machine-moon/tinax
- Repository: https://github.com/machine-moon/tinax

## Overview

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: [`tinax/weights/_safetensors.py`](https://github.com/machine-moon/tinax/blob/master/tinax/weights/_safetensors.py) (Safetensors interchange), [`tinax/checkpointing/_training.py`](https://github.com/machine-moon/tinax/blob/master/tinax/checkpointing/_training.py) (Training checkpoints), [`tinax/weights/_models.py`](https://github.com/machine-moon/tinax/blob/master/tinax/weights/_models.py) (Weight models), [`tinax/weights/_validation.py`](https://github.com/machine-moon/tinax/blob/master/tinax/weights/_validation.py) (Weight validation)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_compute["JAX primitives"]
  node_array["Array conversion and inspection<br/>[_conversion.py]"]
  node_device["Device discovery and validation<br/>[_detect.py]"]
  node_random["Typed RNG keys<br/>[_keys.py]"]
  node_grad["Autodiff transforms<br/>[_transforms.py]"]
  node_grad_validation["Autodiff validation<br/>[_validation.py]"]
  node_jit["Compiled batching<br/>[_batching.py]"]
  node_jit_compile["Trace-budgeted compilation<br/>[_compile.py]"]
  node_jit_validation["Compilation validation<br/>[_validation.py]"]
  node_debug["Debug observation"]
  node_nn_graph["NNX graph copies<br/>[_graph.py]"]
end

subgraph group_input["Data and checkpoints"]
  node_data["Deterministic input pipelines<br/>[_pipelines.py]"]
  node_data_transforms["Batch transforms<br/>[_generic.py]"]
  node_checkpoint["Training checkpoints<br/>[_training.py]"]
  node_checkpoint_ops["Checkpoint operations<br/>[_operations.py]"]
  node_checkpoint_names["Checkpoint naming<br/>[_names.py]"]
end

subgraph group_distributed["Parallel computing"]
  node_parallel_arrays["Distributed arrays<br/>[_arrays.py]"]
  node_mesh["Device meshes<br/>[_mesh.py]"]
  node_shard_map["Manual SPMD mapping<br/>[_shard_map.py]"]
  node_parallel_validation["Layout validation<br/>[_validation.py]"]
end

subgraph group_interchange["Weight interchange"]
  node_weights["Safetensors interchange<br/>[_safetensors.py]"]
  node_weight_files["Weight file handling<br/>[_files.py]"]
  node_weight_formats["Dtype and format rules<br/>[_formats.py]"]
  node_weight_models["Weight models<br/>[_models.py]"]
  node_weight_validation["Weight validation<br/>[_validation.py]"]
end

subgraph group_application["Application boundaries"]
  node_stdlib["CLI and logging boundaries"]
end

node_user(("Library user"))

node_user -->|"converts arrays"| node_array
node_user -->|"inspects devices"| node_device
node_user -->|"derives keys"| node_random
node_user -->|"differentiates functions"| node_grad
node_user -->|"batches functions"| node_jit
node_user -->|"compiles functions"| node_jit_compile
node_user -->|"observes execution"| node_debug
node_user -->|"copies graphs"| node_nn_graph
node_user -->|"builds input pipelines"| node_data
node_user -->|"saves training state"| node_checkpoint
node_user -->|"constructs distributed arrays"| node_parallel_arrays
node_user -->|"configures parallelism"| node_mesh
node_user -->|"maps SPMD work"| node_shard_map
node_user -->|"interchanges weights"| node_weights
node_user -->|"sets app boundaries"| node_stdlib
node_grad -->|"validates arguments"| node_grad_validation
node_jit -->|"validates settings"| node_jit_validation
node_jit_compile -->|"validates settings"| node_jit_validation
node_data -->|"uses batch transform"| node_data_transforms
node_checkpoint -->|"loads and saves"| node_checkpoint_ops
node_checkpoint -->|"resolves names"| node_checkpoint_names
node_checkpoint_ops -->|"validates names"| node_checkpoint_names
node_parallel_arrays -->|"validates layouts"| node_parallel_validation
node_mesh -->|"validates mesh"| node_parallel_validation
node_shard_map -.->|"validates layout"| node_parallel_validation
node_weights -->|"handles files"| node_weight_files
node_weights -->|"uses dtype rules"| node_weight_formats
node_weights -->|"uses data models"| node_weight_models
node_weights -->|"validates inputs"| node_weight_validation
node_weight_models -->|"uses dtype rules"| node_weight_formats
node_weight_models -->|"validates values"| node_weight_validation
node_weight_validation -->|"checks serialization"| node_weight_formats

click node_array "https://github.com/machine-moon/tinax/blob/master/tinax/array/_conversion.py"
click node_device "https://github.com/machine-moon/tinax/blob/master/tinax/device/_detect.py"
click node_random "https://github.com/machine-moon/tinax/blob/master/tinax/random/_keys.py"
click node_grad "https://github.com/machine-moon/tinax/blob/master/tinax/grad/_transforms.py"
click node_grad_validation "https://github.com/machine-moon/tinax/blob/master/tinax/grad/_validation.py"
click node_jit "https://github.com/machine-moon/tinax/blob/master/tinax/jit/_batching.py"
click node_jit_compile "https://github.com/machine-moon/tinax/blob/master/tinax/jit/_compile.py"
click node_jit_validation "https://github.com/machine-moon/tinax/blob/master/tinax/jit/_validation.py"
click node_debug "https://github.com/machine-moon/tinax/tree/master/tinax/debug"
click node_nn_graph "https://github.com/machine-moon/tinax/blob/master/tinax/nn/_graph.py"
click node_data "https://github.com/machine-moon/tinax/blob/master/tinax/data/_pipelines.py"
click node_data_transforms "https://github.com/machine-moon/tinax/blob/master/tinax/data/_generic.py"
click node_checkpoint "https://github.com/machine-moon/tinax/blob/master/tinax/checkpointing/_training.py"
click node_checkpoint_ops "https://github.com/machine-moon/tinax/blob/master/tinax/checkpointing/_operations.py"
click node_checkpoint_names "https://github.com/machine-moon/tinax/blob/master/tinax/checkpointing/_names.py"
click node_parallel_arrays "https://github.com/machine-moon/tinax/blob/master/tinax/parallel/_arrays.py"
click node_mesh "https://github.com/machine-moon/tinax/blob/master/tinax/parallel/_mesh.py"
click node_shard_map "https://github.com/machine-moon/tinax/blob/master/tinax/parallel/_shard_map.py"
click node_parallel_validation "https://github.com/machine-moon/tinax/blob/master/tinax/parallel/_validation.py"
click node_weights "https://github.com/machine-moon/tinax/blob/master/tinax/weights/_safetensors.py"
click node_weight_files "https://github.com/machine-moon/tinax/blob/master/tinax/weights/_files.py"
click node_weight_formats "https://github.com/machine-moon/tinax/blob/master/tinax/weights/_formats.py"
click node_weight_models "https://github.com/machine-moon/tinax/blob/master/tinax/weights/_models.py"
click node_weight_validation "https://github.com/machine-moon/tinax/blob/master/tinax/weights/_validation.py"
click node_stdlib "https://github.com/machine-moon/tinax/tree/master/tinax/stdlib"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_array,node_device,node_random,node_grad,node_grad_validation,node_jit,node_jit_compile,node_jit_validation,node_debug,node_nn_graph,node_user toneBlue
class node_data,node_data_transforms,node_checkpoint,node_checkpoint_ops,node_checkpoint_names toneAmber
class node_parallel_arrays,node_mesh,node_shard_map,node_parallel_validation toneMint
class node_weights,node_weight_files,node_weight_formats,node_weight_models,node_weight_validation toneRose
class node_stdlib toneIndigo
```

## Components

### JAX primitives

- [Array conversion and inspection](https://github.com/machine-moon/tinax/blob/master/tinax/array/_conversion.py): `tinax/array/_conversion.py`
- [Device discovery and validation](https://github.com/machine-moon/tinax/blob/master/tinax/device/_detect.py): `tinax/device/_detect.py`
- [Typed RNG keys](https://github.com/machine-moon/tinax/blob/master/tinax/random/_keys.py): `tinax/random/_keys.py`
- [Autodiff transforms](https://github.com/machine-moon/tinax/blob/master/tinax/grad/_transforms.py): `tinax/grad/_transforms.py`
- [Autodiff validation](https://github.com/machine-moon/tinax/blob/master/tinax/grad/_validation.py): `tinax/grad/_validation.py`
- [Compiled batching](https://github.com/machine-moon/tinax/blob/master/tinax/jit/_batching.py): `tinax/jit/_batching.py`
- [Trace-budgeted compilation](https://github.com/machine-moon/tinax/blob/master/tinax/jit/_compile.py): `tinax/jit/_compile.py`
- [Compilation validation](https://github.com/machine-moon/tinax/blob/master/tinax/jit/_validation.py): `tinax/jit/_validation.py`
- [Debug observation](https://github.com/machine-moon/tinax/tree/master/tinax/debug): `tinax/debug`
- [NNX graph copies](https://github.com/machine-moon/tinax/blob/master/tinax/nn/_graph.py): `tinax/nn/_graph.py`

### Data and checkpoints

- [Deterministic input pipelines](https://github.com/machine-moon/tinax/blob/master/tinax/data/_pipelines.py): `tinax/data/_pipelines.py`
- [Batch transforms](https://github.com/machine-moon/tinax/blob/master/tinax/data/_generic.py): `tinax/data/_generic.py`
- [Training checkpoints](https://github.com/machine-moon/tinax/blob/master/tinax/checkpointing/_training.py): `tinax/checkpointing/_training.py`
- [Checkpoint operations](https://github.com/machine-moon/tinax/blob/master/tinax/checkpointing/_operations.py): `tinax/checkpointing/_operations.py`
- [Checkpoint naming](https://github.com/machine-moon/tinax/blob/master/tinax/checkpointing/_names.py): `tinax/checkpointing/_names.py`

### Parallel computing

- [Distributed arrays](https://github.com/machine-moon/tinax/blob/master/tinax/parallel/_arrays.py): `tinax/parallel/_arrays.py`
- [Device meshes](https://github.com/machine-moon/tinax/blob/master/tinax/parallel/_mesh.py): `tinax/parallel/_mesh.py`
- [Manual SPMD mapping](https://github.com/machine-moon/tinax/blob/master/tinax/parallel/_shard_map.py): `tinax/parallel/_shard_map.py`
- [Layout validation](https://github.com/machine-moon/tinax/blob/master/tinax/parallel/_validation.py): `tinax/parallel/_validation.py`

### Weight interchange

- [Safetensors interchange](https://github.com/machine-moon/tinax/blob/master/tinax/weights/_safetensors.py): `tinax/weights/_safetensors.py`
- [Weight file handling](https://github.com/machine-moon/tinax/blob/master/tinax/weights/_files.py): `tinax/weights/_files.py`
- [Dtype and format rules](https://github.com/machine-moon/tinax/blob/master/tinax/weights/_formats.py): `tinax/weights/_formats.py`
- [Weight models](https://github.com/machine-moon/tinax/blob/master/tinax/weights/_models.py): `tinax/weights/_models.py`
- [Weight validation](https://github.com/machine-moon/tinax/blob/master/tinax/weights/_validation.py): `tinax/weights/_validation.py`

### Application boundaries

- [CLI and logging boundaries](https://github.com/machine-moon/tinax/tree/master/tinax/stdlib): `tinax/stdlib`

### Other components

- **Library user**

## Connections

- Library user → Array conversion and inspection: converts arrays (evidence: `README.md`)
- Library user → Device discovery and validation: inspects devices (evidence: `README.md`)
- Library user → Typed RNG keys: derives keys (evidence: `README.md`)
- Library user → Autodiff transforms: differentiates functions (evidence: `README.md`)
- Library user → Compiled batching: batches functions (evidence: `README.md`)
- Library user → Trace-budgeted compilation: compiles functions (evidence: `README.md`)
- Library user → Debug observation: observes execution (evidence: `README.md`)
- Library user → NNX graph copies: copies graphs (evidence: `README.md`)
- Library user → Deterministic input pipelines: builds input pipelines (evidence: `README.md`)
- Library user → Training checkpoints: saves training state (evidence: `README.md`)
- Library user → Distributed arrays: constructs distributed arrays (evidence: `README.md`)
- Library user → Device meshes: configures parallelism (evidence: `README.md`)
- Library user → Manual SPMD mapping: maps SPMD work (evidence: `README.md`)
- Library user → Safetensors interchange: interchanges weights (evidence: `README.md`)
- Library user → CLI and logging boundaries: sets app boundaries (evidence: `README.md`)
- Autodiff transforms → Autodiff validation: validates arguments (evidence: `tinax/grad/_transforms.py`)
- Compiled batching → Compilation validation: validates settings (evidence: `tinax/jit/_batching.py`)
- Trace-budgeted compilation → Compilation validation: validates settings (evidence: `tinax/jit/_compile.py`)
- Deterministic input pipelines → Batch transforms: uses batch transform (evidence: `tinax/data/_pipelines.py`)
- Training checkpoints → Checkpoint operations: loads and saves (evidence: `tinax/checkpointing/_training.py`)
- Training checkpoints → Checkpoint naming: resolves names (evidence: `tinax/checkpointing/_training.py`)
- Checkpoint operations → Checkpoint naming: validates names (evidence: `tinax/checkpointing/_operations.py`)
- Distributed arrays → Layout validation: validates layouts (evidence: `tinax/parallel/_arrays.py`)
- Device meshes → Layout validation: validates mesh (evidence: `tinax/parallel/_mesh.py`)
- Manual SPMD mapping → Layout validation: validates layout
- Safetensors interchange → Weight file handling: handles files (evidence: `tinax/weights/_safetensors.py`)
- Safetensors interchange → Dtype and format rules: uses dtype rules (evidence: `tinax/weights/_safetensors.py`)
- Safetensors interchange → Weight models: uses data models (evidence: `tinax/weights/_safetensors.py`)
- Safetensors interchange → Weight validation: validates inputs (evidence: `tinax/weights/_safetensors.py`)
- Weight models → Dtype and format rules: uses dtype rules (evidence: `tinax/weights/_models.py`)
- Weight models → Weight validation: validates values (evidence: `tinax/weights/_models.py`)
- Weight validation → Dtype and format rules: checks serialization (evidence: `tinax/weights/_validation.py`)