mlx-loop 🔁

An agentic, closed-loop machine learning experimenter optimized for local execution on Apple Silicon using MLX.

Inspired by the mission to automate the scientific and ML experimental loop, mlx-loop demonstrates a self-correcting ML engineer that designs models, writes training code, executes runs, analyzes logs, and self-patches bugs—all locally on your Mac.

---

Why Apple Silicon & MLX?

Traditional agentic ML workflows (like running agents that train models on cloud VMs) are expensive, high-latency, and require complex orchestration. By leveraging Apple Silicon's unified memory architecture and the lightweight MLX framework:

• Low Latency: Sub-second starting times for training runs compared to spinning up cloud instances.
• Zero Compute Cost: Run iterations locally on your M-series GPU without expensive AWS or GCP bills.
• Privacy-first: Datasets and trained weights stay on local disk.

---

How It Works

graph TD
    A[Task Specification] --> B(Agent: Gemini API)
    B -->|Generates MLX Training Script| C(Executor: Sandbox Process)
    C -->|Runs on M-Series GPU| D{Check Output}
    D -->|Crash or Syntax Error| E[Agent reviews Stderr & updates Code]
    D -->|Metric Not Met <95%| F[Agent tunes Hyperparameters]
    D -->|Success!| G[Save Model Weights model.npz]
    E --> B
    F --> B

Loading

1. Task Definition: Define a target problem (e.g. non-linear spiral classification).
2. Code Generation: The Gemini-powered agent writes a standalone training script using mlx.
3. Subprocess Execution: The script runs locally on Apple Silicon.
4. Log Analysis & Correction:
   • Syntax/Runtime Crashes: If the script crashes (e.g., shape mismatch, missing import), the agent receives the traceback and applies a code correction.
   • Performance Tuning: If it runs successfully but validation accuracy falls below 95%, the agent tunes layers, learning rates, or optimizer configs and runs it again.
5. Convergence: The loop terminates successfully once the validation accuracy threshold is cleared.

---

Directory Structure

mlx-loop/
├── src/
│   └── mlx_loop/
│       ├── agent.py      # LLM reasoning & MLX code-repair logic
│       ├── executor.py   # Subprocess runner & metric parser
│       ├── task.py       # Dataset generator & specification
│       └── cli.py        # Loop orchestrator
├── test_mock.py          # Offline test suite simulating loop iterations
├── pyproject.toml        # Package definition
└── README.md             # This document

---

Setup

1. Clone & Setup Environment:

   git clone https://github.com/yourusername/mlx-loop.git
   cd mlx-loop
   python3 -m venv .venv
   source .venv/bin/activate
   pip install -e .

2. Configure API Key: Create a .env file in the root directory:

   GEMINI_API_KEY=your_gemini_api_key_here

---

Running the Loop

1. Run the Verification Test (Offline / Mocked)

Verify the orchestration and repair loop without hitting your API token quota:

python test_mock.py

2. Run the Live Loop

Kick off the live agentic training run:

python -m mlx_loop.cli

---

Technical Details

• MLX Lazy Evaluation: The agent is guided via a specialized MLX cheatsheet to ensure it calls mlx.core.eval() on model parameters and optimizer state inside training loops, avoiding memory memory leaks or execution graphs that never materialize.
• Subprocess Isolation: Generated scripts are executed in sandboxed subprocesses with a 45-second timeout constraint to prevent infinite loops.