google-deepmind/gemma: a JAX library for running and fine-tuning Gemma weights
Gemma open-weight LLM library, from Google DeepMind
At a glance
- What is it?
- The gemma PyPI package is a JAX implementation for sampling from and fine-tuning Gemma checkpoints, with a ChatSampler that spans Gemma 2, 3, 3n and 4. It is a research-oriented library, not a deployment runtime, and the README leaves several operational questions open.
- Who is it for?
- Adopt this library if you already work in JAX and want to fine-tune or sample from Gemma 2, 3, 3n or 4 weights without leaving Python, using gm.nn, gm.ckpts and gm.text.ChatSampler as shown in the README. Do not adopt it if you need a serving stack with batching, admission control or a stable HTTP contract; the repository documents a library, not a server.
- Can I use it commercially?
- Yes. Apache-2.0 is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 13 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
What the gemma package actually solves
Running an open-weight model usually means choosing between a serving runtime you do not control and a reference implementation you have to port. This repository takes the second path. It is the implementation of the gemma PyPI package, described in the README as a JAX library to use and fine-tune Gemma, where Gemma itself is a family of open-weights models from Google DeepMind based on Gemini research. The audience is narrow and identifiable: people who already have JAX installed for CPU, GPU or TPU work and want model weights to load into that stack rather than into a separate inference server. If your day-to-day is PyTorch, the value proposition is weaker, because the library gives you JAX arrays, JAX compilation and JAX device placement, and nothing else. If your day-to-day is JAX, the library removes the work of reimplementing the architecture and the checkpoint mapping, and gives you a sampler you can call from a notebook.
The ChatSampler API and what it hides
The README's minimal example is the clearest statement of the design. You construct a model with gm.nn.Gemma4_E4B(), load weights with gm.ckpts.load_params(gm.ckpts.CheckpointPath.GEMMA4_E4B_IT), then wrap both in gm.text.ChatSampler with multi_turn=True. Conversation state lives in the sampler object, so the second call takes only the follow-up string and no history. Multi-modal input is expressed inline in the prompt text: the example places <|image|> tokens inside the string and passes images=[image1, image2] as a separate argument. That is a deliberate split. The text carries the position, the argument carries the pixels. The README states the same ChatSampler API works with all Gemma versions (2, 3, 3n, 4), which is the strongest compatibility claim in the material and also the one worth testing first on your own checkpoint. What the example does not show is tokenization, sampling parameters, batch handling or device sharding. Those are left to the documentation and the examples/ folder, which the README says contains additional scripts to fine-tune and sample.
Installation is two steps, and the first one is yours
The README gives an unusually short installation path. Step one is to install JAX for CPU, GPU or TPU by following the instructions on the JAX website. Step two is pip install gemma. The ordering matters: the gemma package does not pin or select your JAX build, so the accelerator support you end up with is whatever you installed first. That is a normal arrangement for a JAX-ecosystem library, but it means a broken CUDA or TPU setup will surface as a JAX error rather than a gemma error, and the README offers no troubleshooting section for it. The system requirements section is the only hardware guidance given: Gemma can run on CPU, GPU and TPU, and for GPU the README recommends 8GB+ RAM on GPU for the 2B checkpoint and 24GB+ RAM on GPU for the 7B checkpoint. Note that these figures are stated for the 2B and 7B checkpoints specifically. The README does not give memory guidance for the Gemma 4 sizes referenced in the code example, so you should treat the 8GB and 24GB numbers as illustrative of scale rather than as a sizing table for the model you intend to load.
Checkpoints are a separate download with separate terms
The code and the weights are distributed differently, and the README is explicit about the split. The repository is Apache-2.0. The weights are not in the repository. The README's Downloading the models section says only: To download the model weights, see our documentation, pointing at the checkpoints page on gemma-llm.readthedocs.io. That page is outside the supplied material, so the exact download mechanism, the authentication requirements and the licence attached to the weights cannot be confirmed here. What can be said is that a checkpoint path constant such as gm.ckpts.CheckpointPath.GEMMA4_E4B_IT implies the library knows where to look, and the checkpoints documentation is where that resolution is defined. For anyone evaluating this for production, the weight licence is the item to read before the code licence, because Apache-2.0 on the Python package tells you nothing about your rights to the parameters.
Where this library is the wrong tool
The README closes with a line that deserves more weight than its position suggests: This is not an official Google product. Combined with the repository contents, that sets expectations. This is a library for using and fine-tuning weights, and the documented surface is a Python API plus Colabs and an examples/ folder. There is no mention of a server, a request queue, continuous batching, or an OpenAI-compatible endpoint. If your requirement is to serve many concurrent users behind an HTTP interface, this is the wrong layer, and you would be writing that layer yourself on top of ChatSampler. A second limitation is version coupling. The README claims one ChatSampler API across Gemma 2, 3, 3n and 4, but the release history shows how fast the surface moves: v3.3.0 in November 2025, then v4.0.0 in May 2026 and v4.0.1 a week later. A major version bump inside roughly six months is a signal to pin your dependency and read the release notes before upgrading, not to track main. Third, the material gives no guidance on distributed inference. Nothing in the README describes sharding a checkpoint across multiple devices, so multi-GPU or multi-host scaling is unverified from what is provided.
How it differs from a Transformers-style implementation
The obvious comparison is Hugging Face Transformers, which also loads Gemma checkpoints and also offers a chat-style generation loop. The difference is the execution model, not the feature list. Transformers is built around PyTorch tensors and eager execution with an ecosystem of quantisation, adapters and serving wrappers attached. This library is built around JAX, which means jit compilation, functional parameter passing and explicit device arrays. In the README example, params is a value you hold and pass into the sampler, not a stateful module attribute. That functional style is what makes the JAX path attractive for fine-tuning experiments where you want to transform or replicate parameters, and it is also what makes it awkward if your surrounding infrastructure assumes torch.nn.Module objects. Choosing between them is mostly a question of which framework your team already debugs at 2am. The README also points to the Gemma ecosystem page at ai.google.dev for other implementations, which is the honest acknowledgement that this is one option among several.
Maintenance cost and the licence boundary
Two maintenance facts are visible. First, the release cadence: v3.3.0, then v4.0.0, then v4.0.1. Patch releases arriving days after a major release usually mean the major release had rough edges, so treating v4.0.0 as the stable target is not supported by the timeline. Second, the test signal: the README carries a Unittests badge pointing at a pytest_and_autopublish workflow, which indicates the project runs a test suite and publishes to PyPI from CI. That is a process signal, not a quality claim, and the material does not include coverage numbers or a support policy. On licensing, the repository is Apache-2.0, which is permissive for the code you would be importing. The weights are governed by whatever the checkpoints documentation specifies, and the README does not restate it. This is not legal advice: read the checkpoint page and any model card linked from the README's technical reports list before you ship weights in a product. The README also invites issues and feedback on GitHub and links contributing guidelines, so the practical upgrade path is to pin gemma in your requirements file, watch the release notes for the next major bump, and keep your checkpoint loading isolated behind the gm.ckpts calls shown in the example so a path or constant change is a one-line fix.
Editorial conclusion
Adopt this library if you already work in JAX and want to fine-tune or sample from Gemma 2, 3, 3n or 4 weights without leaving Python, using gm.nn, gm.ckpts and gm.text.ChatSampler as shown in the README. Do not adopt it if you need a serving stack with batching, admission control or a stable HTTP contract; the repository documents a library, not a server. Before committing, verify the checkpoint download procedure on the checkpoints page, confirm which JAX backend and version your hardware requires, and check whether the model weights carry terms separate from the Apache-2.0 code licence. The README also states plainly that this is not an official Google product, which is the boundary to keep in mind when deciding how much to depend on it.
Community notes