google-deepmind/simply: A JAX Codebase Built for Forking, Not for Configuring

The problem Simply targets: the cost of editing someone else's training loop

Most LLM training repositories make a specific bet. They bet that you will spend more time configuring than writing model code, so they grow a config system, a registry of components, and a plugin boundary between the framework and your idea. That bet pays off for production pipelines. It pays off badly for a researcher who wants to change the training loss on a Tuesday afternoon and needs to trace that loss through four levels of indirection first.

Simply inverts the bet. The README states the goal directly: minimizing the time to implement new ideas such as an optimizer, a training loss, or an RL algorithm, by humans and by AI agents. The named audience is frontier LLM research, and the second audience is explicit: an AI agent that can read the code, propose ideas, run experiments, and iterate, either autonomously or under human guidance. That second audience shapes the design more than the first. An agent has no institutional memory of your codebase and no patience for abstraction layers, so a repository that a model can read end to end in one context window is worth more to it than a repository with a clean plugin API.

The consequence is a codebase that is short by construction rather than short by accident. The README's own framing is that if you learn JAX, you are ready to read and hack the code. That is a real constraint on who can use this: the prerequisite is not PyTorch familiarity, it is JAX familiarity, including jit and scan semantics, because the codebase does not hide them.

What the repository layout tells you about the architecture

The layout is flat and legible. There is a simply package holding the training code, a simply/serving/ directory for a gRPC serving stack, a simply/agent/ directory for a built-in agent harness, an amplio/ directory for a vendored Go agent harness, and a setup/ directory with asset and proto generation scripts. Documentation lives in docs/ and in top-level markdown files such as gcloud_quickstart.md.

The dependency list is the architecture in miniature: JAX for model and training, Orbax for checkpoint management, Grain for the data pipeline. Three dependencies, each with a distinct job. There is no experiment tracking service, no configuration framework, and no model registry in that list. Checkpointing is delegated to Orbax rather than hand-rolled, which is the right call: checkpoint resumption is exactly the kind of code that is boring to write and expensive to get wrong.

The serving layer is worth noting because it is where the abstraction budget is spent. According to the README, simply/serving/ talks gRPC and its Python stubs are generated from .proto files rather than checked in. That is a deliberate choice with a cost attached, covered below. The agent directories are the unusual part of the layout. Two agent harnesses ship in the same repository as the training code, which means the repository is not only a training stack but also a substrate for automated experimentation on itself.

Getting a first run: the exact commands and where they write

The README gives a local debug command that runs a small language model test config and writes to a temporary directory:

EXP=simply_local_test_1; rm -rf /tmp/${EXP}; python -m simply.main --experiment_config lm_test --experiment_dir /tmp/${EXP} --alsologtostderr

A second variant disables jit and scan so you can print arrays like ordinary Python values:

export JAX_DISABLE_JIT=True; EXP=simply_local_test_1; rm -rf /tmp/${EXP}; python -m simply.main --experiment_config lm_no_scan_test --experiment_dir /tmp/${EXP} --alsologtostderr

That second command is the most useful thing in the README for anyone evaluating the project. Disabling scan means the per-step loop unrolls, so intermediate arrays become inspectable. It also means the run is slow and memory-hungry, which is the point: it is a debugging mode, not a training mode.

Installation is split by accelerator. The README advises installing JAX per environment (pip install -U jax for CPU, pip install -U "jax[cuda13]" for GPU, pip install -U "jax[tpu]" for TPU), then pip install . for the package itself. Optional extras are named: .[tfds] for TensorFlow Datasets, .[math-eval] for simply/utils/math_eval.py, .[dev] for pytest, .[serving] for the gRPC stubs, and .[agent] for the built-in harness.

Assets come from HuggingFace. python setup/setup_assets.py downloads both models and datasets, with --models-only and --datasets-only flags, landing in ~/.cache/simply/models/ and ~/.cache/simply/datasets/. Both locations are overridable via --models-dir and --datasets-dir or the SIMPLY_MODELS and SIMPLY_DATASETS environment variables. The README adds a caveat in parentheses that matters for planning: only a few datasets and models are included for testing, with more promised later. The uv path collapses dependency install and asset download into uv sync followed by uv run python setup/setup_assets.py.

Two agent harnesses, two very different maintenance surfaces

The built-in harness in simply/agent/ is documented with a Bash tool and context management, aimed at autonomous long-running research tasks. Its quick test is short:

pip install ".[agent]" python -m simply.agent.main --task_file=simply/agent/example_tasks/code_stats.md --env="Local:." --llm="LiteLLM:vertex_ai/gemini-2.5-pro"

The task file, the environment string, and the LLM string are all positional-style flags, so swapping the model provider is a one-argument change. The example task is a code statistics task, which is a reasonable smoke test because it needs filesystem access but no accelerator.

Amplio is the heavier option and it is vendored, meaning its source lives under amplio/ in this repository rather than being pulled in as a dependency. The README describes it as a Go agent harness with generic tools (shell, file edit, sub-agent spawn, inter-agent messaging), DB-first persistence so a run resumes after a crash, and a web UI. Building it requires Go and Node.js 22 or later, per amplio/go.mod and the README: cd amplio && make build, then ./amplio serve, which prints a URL with an access token. Configuration is read from <data-dir>/config.toml under the keys system_llm_hq and system_llm_fast.

Here is the trade-off to weigh. Vendoring gives you a harness that is pinned to this repository and cannot drift out from under you. It also means a Go toolchain and a Node.js 22 toolchain are now part of your research environment, and that the harness's upgrade cadence is this repository's commit cadence, not an independent project's release cycle. If your work is optimizer and loss experiments on a single TPU slice, the built-in Python harness is the lighter path. Amplio's DB-first persistence is the feature that justifies its weight, and it only justifies it if your runs are long enough that a crash mid-run is a real loss.

Where Simply is the wrong tool

The most concrete limitation is stated in the README itself: only a few datasets and models are included for testing. If your research depends on a specific pretrained checkpoint or a particular corpus, setup/setup_assets.py may not have it, and you will be writing your own data and checkpoint plumbing into a codebase whose entire value proposition is that it does not have much plumbing. That is a real mismatch, not a minor gap.

The second constraint is the generated serving stubs. The README states that the Python stubs for simply/serving/ are generated from .proto files rather than checked in, and that you must run python setup/gen_protos.py once after installing with the .[serving] extra. Anyone who clones the repository and imports the serving package without that step will hit a failure that has nothing to do with their research idea. It is a one-line fix, but it is a setup step that a fresh checkout does not satisfy.

The third is the absence of retrieved releases. There are no tagged versions in the material supplied, so there is no versioned upgrade path to reason about. Combined with the fork-and-hack framing, this points to a workflow where you clone or fork and track main, which means you own your divergence. If your team needs pinned dependency versions and a changelog to diff against before upgrading, this repository does not currently offer that in the material available.

Finally, the codebase is minimal by intent. Every convenience you might expect from a general training framework (a config schema, a component registry, multi-model support, built-in evaluation suites) is either absent or thin. The .[math-eval] extra covers simply/utils/math_eval.py, which suggests evaluation is present but narrow. If you need breadth, minimalism is a liability rather than a feature.

How this differs from a general-purpose training framework

Take a framework like Hugging Face's Trainer or PyTorch Lightning as the comparison point. Both are built so that a user expresses intent through configuration and callbacks while the framework owns the training loop. You get a stable interface, model and dataset breadth, and an upgrade path where your code keeps working across versions. You pay for it in indirection: changing the loss function means finding the right hook, and reading the whole thing is not feasible in one sitting.

Simply makes the opposite trade. The training loop is yours to edit, the dependencies are three (JAX, Orbax, Grain), and the README's claim is that knowing JAX is sufficient preparation. The cost is that there is no interface stability to lean on. When you fork, you own the fork. When upstream changes the optimizer interface in simply/utils/optimizers.py, the README's own agent prompt tells the agent to read that file to understand the interface, which is a hint that the interface is defined by the code rather than by documentation.

There is a second, less obvious difference. General frameworks are designed for a human who reads documentation. Simply is designed for a reader that reads source. The README's example agent prompt instructs the agent to design and benchmark new optimizers, run the Adam baseline on lm_test, propose novel optimizers, run three experiments per iteration, stop after 15 experiments or 10 proposed optimizers, and write a report to /tmp/optimizer_report.md. That prompt only works if the repository is small enough to fit in a context window and explicit enough that an agent can find the optimizer interface without a guide. A large framework with a plugin system would defeat the premise.

Licence and the cost of staying current

Simply is Apache-2.0. That is a permissive licence that permits commercial use and modification, and it includes a patent grant. It also carries the usual obligations: you must retain the licence and attribution notices in redistributed copies, and state significant changes you made. This is a description of the licence text, not legal advice; if you are redistributing a modified Simply inside a product, have counsel read the NOTICE handling and the modification-statement requirement.

One licence detail worth flagging because of the repository layout: amplio/ is vendored under this repository, so it ships under whatever terms the repository states unless a separate licence file exists inside that directory. The material supplied does not show a separate licence for amplio/, and the top-level licence is Apache-2.0. If you plan to redistribute the harness, check amplio/ for its own licence file before assuming the top-level terms apply.

Upgrade cost is the harder question. With no releases retrieved and a fork-first workflow, the practical maintenance model is: you pull from main, resolve conflicts against your own edits, and re-run the local test configs to see whether anything broke. The two commands in the README (lm_test and lm_no_scan_test) are the natural regression check, since they exercise the training path end to end without needing a TPU. Budget for that conflict resolution as a recurring cost, not a one-time setup fee. The more of the training loop you rewrite, the higher it gets.

Editorial conclusion

Adopt Simply if your research output is a code diff (a new optimizer, loss, or RL algorithm) and you are already fluent in JAX, since the README's own pitch is that you learn JAX and are then ready to read the code. Do not adopt it if you need stable configuration surfaces, a broad model zoo, or a documented upgrade path: the README states that only a few models and datasets are included for testing, and no releases were retrieved. Before committing, run the lm_test config locally, confirm setup_assets.py populates ~/.cache/simply/models/ and ~/.cache/simply/datasets/, and check whether your workflow needs the serving protos generated by setup/gen_protos.py, since those stubs are not checked in.

Official sources

1. google-deepmind/simply on GitHub
2. Issues
3. License: Apache-2.0
4. README