The gemma Python package: a JAX library for running and fine-tuning Gemma 4

What the gemma package actually is, and who it is not for

Gemma is a family of open-weights language models from Google DeepMind, based on Gemini research. The repository at google-deepmind/gemma is not the model and not an application. It is the implementation of the gemma PyPI package, described in the README as "a JAX library to use and fine-tune Gemma." That sentence defines the audience precisely: people who write Python, who are comfortable with JAX, and who want the model graph, the parameter loading and the training step under their own control.

If you want a chat window, a downloaded binary or a local server you start with one command, this is the wrong repository. The README's own framing puts the library in a position below those tools. It links out to a separate Gemma ecosystem page for other implementations. The package assumes you are assembling the pieces yourself: pick a checkpoint, instantiate a model class, load parameters, then decide what to do with them.

The repository layout supports that reading. There is a gemma/ package directory, a docs/ folder, a colabs/ folder and an examples/ directory holding scripts such as classification.py, dpo.py, lora.py, multimodal.py, npo.py, seq2seq.py and sharding.py. Those filenames are the real map of what the library is for. Classification, preference optimisation, LoRA adapters, sequence-to-sequence work and sharding are all first-class concerns here. General chat is one example among several.

How the JAX model, checkpoint and sampler fit together

The mechanism is visible in the README's minimal example. You construct a model object from a class in gm.nn, load a parameter tree from a checkpoint path in gm.ckpts, and hand both to a sampler in gm.text. The sampler owns the generation loop. Inference is a pure function of model plus parameters, which is the JAX idiom: the model object describes the computation, the parameter tree holds the weights, and nothing is stored inside the model as mutable state.

The checkpoint enum is worth noting. The example uses gm.ckpts.CheckpointPath.GEMMA4_E4B_IT, and the naming carries information: E4B identifies the variant, and the IT suffix marks the instruction-tuned checkpoint rather than the base one. The README states that the same ChatSampler API works with all Gemma versions, 2, 3, 3n and 4. That is a deliberate compatibility promise, and it means the model class you instantiate is the main thing that changes between versions, not the calling convention.

Multi-modality is handled at the prompt level rather than through a separate pipeline. The example prompt embeds an <|image|> token where an image belongs, and the chat call passes an images list positionally alongside the text. Multi-turn state is a flag on the sampler, not something you re-feed: with multi_turn=True, the second call takes only the new user turn and the sampler carries the history. That is a small design decision with a real consequence, since conversation state lives in the sampler object and not in your prompt string.

The dependency list in pyproject.toml confirms the shape of the stack. Alongside jax and jaxtyping, the package pulls in kauldron, etils, einops, grain on non-Windows platforms, and hackable-diffusion from a Git URL. Kauldron is a training framework, which is consistent with the fine-tuning examples being first-class rather than bolted on.

Installing gemma and holding a first multi-modal conversation

Installation is two steps, and the order matters. The README says to install JAX for CPU, GPU or TPU first, following the JAX website instructions, and only then to install the package. The reason is that JAX's accelerator support is chosen at install time through platform-specific wheels; installing gemma first does not pick the right backend for you.

pip install gemma

That is the entire published install command. Note the Python floor: pyproject.toml sets requires-python to ">=3.12", so an older interpreter will fail at resolve time rather than at import time.

The README's minimal example is the shortest path to something real. It builds a Gemma 4 model, loads the instruction-tuned E4B parameters, and creates a chat sampler configured for multi-turn use.

from gemma import gm

model = gm.nn.Gemma4_E4B()
params = gm.ckpts.load_params(gm.ckpts.CheckpointPath.GEMMA4_E4B_IT)

sampler = gm.text.ChatSampler(
    model=model,
    params=params,
    multi_turn=True,
)

What you should see is a sampler object ready to accept text. The first call takes a prompt, and images are passed as a list. The README's example prompt places two <|image|> tokens in the text and asks the model to compare the two images.

out0 = sampler.chat(prompt, images=[image1, image2])
out1 = sampler.chat('What about the other image ?')

The second call passes no images and no history, because multi_turn=True already retained them. If that second answer ignores the images, the likely cause is that multi_turn was left at its default.

One gap: the README does not show where the weights come from. It says only to see the documentation for downloading model weights, so a first run depends on the checkpoints page at gemma-llm.readthedocs.io rather than on anything in the README itself.

Where the library pushes work back onto you

The honest limitation is resource planning, and the README is unusually direct about it. It recommends 8GB or more of GPU RAM for the 2B checkpoint and 24GB or more for the 7B checkpoint. Those numbers describe the weights, not the full training footprint. Activations, optimiser state and gradients all sit on top of that during fine-tuning, so the examples in examples/lora.py and examples/dpo.py are not interchangeable with the sampling path in memory terms. LoRA exists precisely because full fine-tuning does not fit in the same envelope.

Platform support is uneven in a way that is easy to miss. The grain dependency in pyproject.toml carries a marker restricting it to non-Windows platforms, so the data pipeline the library depends on is not installed on Windows at all. The README does not discuss Windows behaviour, and it does not document rollback or a supported downgrade path between the 3.x and 4.x lines. If you are pinned to an older release, the CHANGELOG.md file is the place to check, not the README.

There is also a licensing distinction that has nothing to do with the code licence. The repository is Apache-2.0, and pyproject.toml declares that classifier. That covers the library source. The model weights are governed separately by the Gemma terms linked from the model documentation, and the README's closing line, that this is not an official Google product, is a support-scope statement rather than a legal one. Treating the Apache-2.0 file as permission to do anything with the weights is a misreading.

Finally, the dependency on hackable-diffusion is pulled from a Git URL rather than a released version. That is a reproducibility hazard in any environment that resolves dependencies from a lock file or a private mirror.

How this differs from running Gemma through a local inference server

The practical alternative for most people searching for Gemma is not another Python library. It is a local inference runtime such as the one the Gemma ecosystem page points to, where you download a quantised model file and talk to it over an HTTP endpoint or a bundled interface. The difference in approach is not cosmetic.

A server runtime owns the whole stack: quantisation, memory mapping, tokenisation, a request queue and often an OpenAI-compatible API. You get a running process and you send it strings. The gemma package gives you none of that. It gives you a model definition that composes with JAX transformations, which is what makes sharding across devices, gradient computation and custom training loops possible. You cannot fine-tune through a chat server, and you cannot inspect per-layer activations through one either.

The trade is operational. With a server runtime, upgrading means pulling a new model file. With this package, upgrading means tracking a Python library whose major version moved from 3.3.0 to 4.0.0 to 4.0.1 within roughly six months, and whose model classes are versioned in the API itself, as Gemma4_E4B shows. That is a normal pace for a research-adjacent library and a fast one for a production dependency.

There is a middle path worth naming: use the server runtime for inference and this package only when you need gradients. Nothing in the README prevents that, and it avoids carrying a JAX and Kauldron install for workloads that only ever call chat().

Maintenance, upgrades and what the licence does not cover

The repository is not archived, and the last push was on 2026-09-03. The release history is short and recent: v3.3.0 in November 2025, then v4.0.0 and v4.0.1 in May 2026. A major version bump between those points means the 3.x to 4.x step is the one to plan for, and the README does not describe a migration path. The CHANGELOG.md at the repository root is the only place that would carry one.

Upgrade cost concentrates in two places. The first is the model class names, since the API surfaces the model generation directly, as Gemma4_E4B demonstrates. The second is the dependency set, which includes kauldron with a minimum version constraint and hackable-diffusion from a Git URL. A Git-sourced dependency means your build is only as reproducible as the commit that URL resolves to at install time.

On licensing, the split is straightforward to state and not to interpret. The library is Apache-2.0, per both the LICENSE file and the classifier in pyproject.toml. That permits commercial use of the code under the usual terms. The weights are a separate artefact with separate terms, and this article is not the place to read them for you. If your use case depends on redistribution or on specific commercial rights, read the Gemma model terms directly rather than inferring them from the repository licence.

Editorial conclusion

Adopt the gemma package if you already work in JAX or Kauldron and want the model definition, the checkpoint loader and the training loop in one Python codebase, or if you need to fine-tune rather than just prompt. Do not adopt it if what you actually want is a desktop chat client or a one-line local server; the README points at Ollama-style tooling only indirectly, through the broader Gemma ecosystem link, and this package is not that. Before committing, verify three things: that your Python is 3.12 or newer, that your JAX install matches the accelerator you own, and that you can reach the checkpoint download page, since the README defers weight acquisition to the documentation rather than shipping a downloader inline.

Frequently asked questions

Official sources

1. google-deepmind/gemma on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases