recurrentgemma

Installation

Using Poetry

RecurrentGemma uses Poetry for dependency management.

To install dependencies for the full project:

• Checkout the code.
• poetry install -E full to create a virtual environment with all dependencies.
• poetry shell to activate the created virtual environment.

If you only need to install a subset of dependencies use one of the alternative library-specific commands below.

Using pip

If you want to use pip instead of Poetry, then create a virtual environment (run python -m venv recurrentgemma-demo and . recurrentgemma-demo/bin/activate) and:

• Checkout the code.
• pip install .[full]

Installing library-specific packages

JAX

To install dependencies only for the JAX pathway use: poetry install -E jax or (pip install .[jax]).

PyTorch

To install dependencies only for the PyTorch pathway use: poetry install -E torch (or pip install .[torch]).

Tests

To install dependencies required for running unit tests use: poetry install -E test (or pip install .[test])

---

System Requirements

RecurrentGemma code can run on CPU, GPU or TPU. The code has been optimized for running on TPU using the Flax implementation, which contains a low level Pallas kernel to perform the linear scan in the recurrent layers.

---

License

Copyright 2024 DeepMind Technologies Limited

This code is licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0.

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an AS IS BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.