# ai-hypercomputer/recml architecture

> Architecture diagram of the GitHub repository ai-hypercomputer/recml, 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/ai-hypercomputer/recml
- Repository: https://github.com/ai-hypercomputer/recml

## Overview

RecML is a recommender-modeling library focused on scalable Keras/JAX workloads, especially Cloud TPU and SparseCore. The principal workflow is to prepare and batch recommendation data, apply a recommendation model, train and evaluate it with a framework-specific trainer, then retain checkpoints or use a model for inference. The README also documents reusable layers, embedding support, and standardized metrics. Some connections are documented capabilities rather than verified calls: sampled excerpts confirm the JAX trainer’s dataset, evaluation, and checkpoint responsibilities, but not every example’s wiring.

Main entry points: [`recml/core/training/jax_trainer.py`](https://github.com/ai-hypercomputer/recml/blob/main/recml/core/training/jax_trainer.py) (JAX trainer), [`recml/core/training`](https://github.com/ai-hypercomputer/recml/tree/main/recml/core/training) (Trainer APIs), [`recml/inference/models/jax/DLRM_DCNv2`](https://github.com/ai-hypercomputer/recml/tree/main/recml/inference/models/jax/DLRM_DCNv2) (DLRM models), [`recml/layers/linen/sparsecore.py`](https://github.com/ai-hypercomputer/recml/blob/main/recml/layers/linen/sparsecore.py) (SparseCore embeddings)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data preparation"]
  node_preprocess["Preprocessing<br/>[preprocessing.py]"]
  node_dataset["Dataset factory"]
  node_iterator["Data iterator<br/>[iterator.py]"]
end

subgraph group_models["Recommendation models"]
  node_sequential["Sequential models"]
  node_hstu["HSTU models"]
  node_dlrm["DLRM models"]
end

subgraph group_training["Training and evaluation"]
  node_trainers["Trainer APIs"]
  node_jaxtrainer["JAX trainer<br/>[jax_trainer.py]"]
  node_kerastrainer["Keras trainer<br/>[keras_trainer.py]"]
  node_partitioning["SPMD partitioning<br/>[partitioning.py]"]
  node_checkpoint[("Checkpoint files")]
end

subgraph group_components["Reusable components"]
  node_metrics["Metrics"]
  node_layers["Recommendation layers"]
  node_sparsecore["SparseCore embeddings<br/>[sparsecore.py]"]
  node_embeddingops["Embedding operations<br/>[embedding_ops.py]"]
end

subgraph group_serving["Inference"]
  node_inference["DLRM inference<br/>[dlrm_main.py]"]
end

node_practitioner(("Practitioner"))
node_records["Input records"]
node_tpu(("Cloud TPU"))

node_practitioner -->|"provides data"| node_records
node_records -->|"prepares features"| node_preprocess
node_preprocess -->|"feeds examples"| node_dataset
node_dataset -->|"batches data"| node_iterator
node_practitioner -->|"configures training"| node_trainers
node_trainers -->|"selects JAX"| node_jaxtrainer
node_trainers -->|"selects Keras"| node_kerastrainer
node_jaxtrainer -->|"consumes datasets"| node_iterator
node_jaxtrainer -->|"saves state"| node_checkpoint
node_jaxtrainer -->|"reports evaluation"| node_metrics
node_kerastrainer -->|"evaluates with"| node_metrics
node_jaxtrainer -->|"uses partitioning"| node_partitioning
node_practitioner -->|"chooses model"| node_sequential
node_practitioner -->|"chooses model"| node_hstu
node_practitioner -->|"chooses model"| node_dlrm
node_sequential -->|"uses building blocks"| node_layers
node_hstu -->|"uses building blocks"| node_layers
node_dlrm -.->|"uses embeddings"| node_sparsecore
node_sparsecore -->|"performs lookup"| node_embeddingops
node_tpu -.->|"accelerates lookup"| node_sparsecore
node_practitioner -->|"runs inference"| node_inference
node_inference -->|"loads model"| node_dlrm

click node_preprocess "https://github.com/ai-hypercomputer/recml/blob/main/recml/core/data/preprocessing.py"
click node_dataset "https://github.com/ai-hypercomputer/recml/blob/main/recml/core/data/tf_dataset_factory.py"
click node_iterator "https://github.com/ai-hypercomputer/recml/blob/main/recml/core/data/iterator.py"
click node_sequential "https://github.com/ai-hypercomputer/recml/tree/main/recml/layers/keras"
click node_hstu "https://github.com/ai-hypercomputer/recml/tree/main/recml/examples/DLRM_HSTU"
click node_dlrm "https://github.com/ai-hypercomputer/recml/tree/main/recml/inference/models/jax/DLRM_DCNv2"
click node_trainers "https://github.com/ai-hypercomputer/recml/tree/main/recml/core/training"
click node_jaxtrainer "https://github.com/ai-hypercomputer/recml/blob/main/recml/core/training/jax_trainer.py"
click node_kerastrainer "https://github.com/ai-hypercomputer/recml/blob/main/recml/core/training/keras_trainer.py"
click node_partitioning "https://github.com/ai-hypercomputer/recml/blob/main/recml/core/training/partitioning.py"
click node_metrics "https://github.com/ai-hypercomputer/recml/tree/main/recml/core/metrics"
click node_layers "https://github.com/ai-hypercomputer/recml/tree/main/recml/layers"
click node_sparsecore "https://github.com/ai-hypercomputer/recml/blob/main/recml/layers/linen/sparsecore.py"
click node_embeddingops "https://github.com/ai-hypercomputer/recml/blob/main/recml/core/ops/embedding_ops.py"
click node_inference "https://github.com/ai-hypercomputer/recml/blob/main/recml/inference/models/jax/DLRM_DCNv2/dlrm_main.py"

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_preprocess,node_dataset,node_iterator toneBlue
class node_sequential,node_hstu,node_dlrm toneAmber
class node_trainers,node_jaxtrainer,node_kerastrainer,node_partitioning,node_checkpoint toneMint
class node_metrics,node_layers,node_sparsecore,node_embeddingops toneRose
class node_inference,node_practitioner,node_records,node_tpu toneIndigo
```

## Components

### Data preparation

- [Preprocessing](https://github.com/ai-hypercomputer/recml/blob/main/recml/core/data/preprocessing.py): `recml/core/data/preprocessing.py`
- [Dataset factory](https://github.com/ai-hypercomputer/recml/blob/main/recml/core/data/tf_dataset_factory.py): `recml/core/data/tf_dataset_factory.py`
- [Data iterator](https://github.com/ai-hypercomputer/recml/blob/main/recml/core/data/iterator.py): `recml/core/data/iterator.py`

### Recommendation models

- [Sequential models](https://github.com/ai-hypercomputer/recml/tree/main/recml/layers/keras): `recml/layers/keras`
- [HSTU models](https://github.com/ai-hypercomputer/recml/tree/main/recml/examples/DLRM_HSTU): `recml/examples/DLRM_HSTU`
- [DLRM models](https://github.com/ai-hypercomputer/recml/tree/main/recml/inference/models/jax/DLRM_DCNv2): `recml/inference/models/jax/DLRM_DCNv2`

### Training and evaluation

- [Trainer APIs](https://github.com/ai-hypercomputer/recml/tree/main/recml/core/training): `recml/core/training`
- [JAX trainer](https://github.com/ai-hypercomputer/recml/blob/main/recml/core/training/jax_trainer.py): `recml/core/training/jax_trainer.py`
- [Keras trainer](https://github.com/ai-hypercomputer/recml/blob/main/recml/core/training/keras_trainer.py): `recml/core/training/keras_trainer.py`
- [SPMD partitioning](https://github.com/ai-hypercomputer/recml/blob/main/recml/core/training/partitioning.py): `recml/core/training/partitioning.py`
- **Checkpoint files**

### Reusable components

- [Metrics](https://github.com/ai-hypercomputer/recml/tree/main/recml/core/metrics): `recml/core/metrics`
- [Recommendation layers](https://github.com/ai-hypercomputer/recml/tree/main/recml/layers): `recml/layers`
- [SparseCore embeddings](https://github.com/ai-hypercomputer/recml/blob/main/recml/layers/linen/sparsecore.py): `recml/layers/linen/sparsecore.py`
- [Embedding operations](https://github.com/ai-hypercomputer/recml/blob/main/recml/core/ops/embedding_ops.py): `recml/core/ops/embedding_ops.py`

### Inference

- [DLRM inference](https://github.com/ai-hypercomputer/recml/blob/main/recml/inference/models/jax/DLRM_DCNv2/dlrm_main.py): `recml/inference/models/jax/DLRM_DCNv2/dlrm_main.py`

### Other components

- **Practitioner**
- **Input records**
- **Cloud TPU**

## Connections

- Practitioner → Input records: provides data
- Input records → Preprocessing: prepares features
- Preprocessing → Dataset factory: feeds examples
- Dataset factory → Data iterator: batches data
- Practitioner → Trainer APIs: configures training
- Trainer APIs → JAX trainer: selects JAX
- Trainer APIs → Keras trainer: selects Keras
- JAX trainer → Data iterator: consumes datasets
- JAX trainer → Checkpoint files: saves state
- JAX trainer → Metrics: reports evaluation
- Keras trainer → Metrics: evaluates with
- JAX trainer → SPMD partitioning: uses partitioning
- Practitioner → Sequential models: chooses model
- Practitioner → HSTU models: chooses model
- Practitioner → DLRM models: chooses model
- Sequential models → Recommendation layers: uses building blocks
- HSTU models → Recommendation layers: uses building blocks
- DLRM models → SparseCore embeddings: uses embeddings
- SparseCore embeddings → Embedding operations: performs lookup
- Cloud TPU → SparseCore embeddings: accelerates lookup
- Practitioner → DLRM inference: runs inference
- DLRM inference → DLRM models: loads model