chore: fix spelling typos in comments and docstrings

examples/classification.py

@@ -96,7 +96,7 @@ def _make_dataset(training: bool) -> kd.data.Pipeline:
           gm.data.FormatText(
               key=_INPUT_FIELD,
               template="""<start_of_turn>user
-              Please classify whether the following sentence is grammaticaly correct, please answer only with Yes or No.
+              Please classify whether the following sentence is grammatically correct, please answer only with Yes or No.
               Sentence: {text}<end_of_turn>
               <start_of_turn>model""",
           ),
@@ -113,7 +113,7 @@ def _make_dataset(training: bool) -> kd.data.Pipeline:
           gm.data.MapInts(
               key=_LABEL_FIELD,
               # Rather than predicting the token 0 and 1, we are using the
-              # token 1294 and 3553 which respectivelly correspond to "No" and
+              # token 1294 and 3553 which respectively correspond to "No" and
               # "Yes". We do this because those token already contain semantic
               # information, so even zero-shot prediction without any
               # finetuning has better than random performances.

examples/dpo.py

@@ -14,7 +14,7 @@
 
 r"""DPO Example.
 
-DPO works by running two answers (one prefered and one rejected) into both
+DPO works by running two answers (one preferred and one rejected) into both
 the reference model and the model to finetune. Then the DPO loss is used to
 increase the likelihood of generating the preferred answer.
 

gemma/gm/ckpts/_checkpoint.py

@@ -49,7 +49,7 @@
 class LoadCheckpoint(kd.ckpts.InitTransform):
   """Loads weights from a Gemma checkpoint.
 
-  Note: The checpoint only contains the Gemma transformer weights, not the
+  Note: The checkpoint only contains the Gemma transformer weights, not the
   step, optimizer state,... Use `kd.ckpts.PartialKauldronLoader` to load
   the state from a Kauldron checkpoint.
 
@@ -292,7 +292,7 @@ def load_params(
   # TODO(epot): Better API. Currently this do not quantize the weights, but
   # just refactor the params to the QAT structure.
   # Eventually quantize the params. Note: It would be better to do this
-  # while the weights are loaded, so restore do not use unecessary memory.
+  # while the weights are loaded, so restore do not use unnecessary memory.
   if quantize:
     output = _quantization.convert_to_qat_checkpoint(output)
 

gemma/gm/ckpts/_lora.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Utils for LoRA checkpoint managment."""
+"""Utils for LoRA checkpoint management."""
 
 from __future__ import annotations
 

gemma/gm/ckpts/_quantization.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Utils for LoRA checkpoint managment."""
+"""Utils for LoRA checkpoint management."""
 
 from __future__ import annotations
 
@@ -53,7 +53,7 @@ def convert_leaf(data: Any) -> Any:
       data = quantize_leaf(data, checkpoint_kernel_key)
     # This hack is required because the FeedForward layer call two different
     # Einsum with using `nn.share_scope`, so the two wrappers need a different
-    # name. Weights are not stored under any kernel key and are isntead under
+    # name. Weights are not stored under any kernel key and are instead under
     # the following names.
     if 'gating_einsum' in data or 'linear' in data:
       data = quantize_leaf(data, 'gating_einsum')

gemma/gm/data/_functional.py

@@ -116,7 +116,7 @@ def make_seq2seq_fields(
   prompt = [10, 11, 12, 13],
   response = [20, 21, 1],  # Here, response ends with EOS token.
 
-  # Ouptut:
+  # Output:
   out.input =       [10, 11, 12, 13, 20, 21],
   out.target =      [11, 12, 13, 20, 21,  1],
   out.target_mask = [ 0,  0,  0,  1,  1,  1],

gemma/gm/data/_tasks.py

@@ -80,7 +80,7 @@ class Seq2SeqTask(grain.MapTransform):
       'prompt': 'Hello! I would love to visit France.',
       'response': 'Bonjour ! J'adorerais visiter la France.',
   }
-  # Ouptut:
+  # Output:
   {
       'input': i32['max_length'],
       'target': i32['max_length 1'],

gemma/gm/data/_transforms.py

@@ -133,7 +133,7 @@ class AddSeq2SeqFields(grain.MapTransform):
       'prompt': [10, 11, 12, 13],
       'response': [20, 21, 1],  # Here, response ends with EOS token.
   }
-  # Ouptut:
+  # Output:
   {
       'input':       [10, 11, 12, 13, 20, 21],
       'target':      [11, 12, 13, 20, 21,  1],

gemma/gm/nn/_transformer.py

@@ -107,7 +107,7 @@ class Transformer(nn.Module):
   attention_mask: kontext.Key | None = None
 
   config: _config.TransformerConfig
-  # Model info to specifiy the tokenizer version and default checkpoint.
+  # Model info to specify the tokenizer version and default checkpoint.
   INFO: ClassVar[ModelInfo] = ModelInfo()
 
   def __post_init__(self):
@@ -178,7 +178,7 @@ def vision_encoder(self) -> gemma_vision.SigLiPFromPatches | None:
           'return_hidden_states',
       ),
   )
-  # The function accepts/returns aribtrary batch shape, but inside the
+  # The function accepts/returns arbitrary batch shape, but inside the
   # function, the batch dimension is flattened to a single dimension.
   @_jax_utils.flatten_unflatten_batch_dim()
   @typechecked

gemma/gm/nn/gemma3n/_transformer.py

@@ -108,7 +108,7 @@ class Gemma3nTransformer(_transformer.Transformer):
   images: kontext.Key | None = None
 
   config: _config.TransformerConfig
-  # Model info to specifiy the tokenizer version and default checkpoint.
+  # Model info to specify the tokenizer version and default checkpoint.
   INFO: ClassVar[ModelInfo] = ModelInfo()
 
   def __post_init__(self):
@@ -217,7 +217,7 @@ def vision_encoder(self) -> gemma_vision.SigLiPFromPatches | None:
           'return_hidden_states',
       ),
   )
-  # The function accepts/returns aribtrary batch shape, but inside the
+  # The function accepts/returns arbitrary batch shape, but inside the
   # function, the batch dimension is flattened to a single dimension.
   @_jax_utils.flatten_unflatten_batch_dim()
   @typechecked

gemma/gm/nn/gemma4/_transformer.py

@@ -128,7 +128,7 @@ class Transformer(nn.Module):
   images: kontext.Key | None = None
 
   config: _config.TransformerConfig
-  # Model info to specifiy the tokenizer version and default checkpoint.
+  # Model info to specify the tokenizer version and default checkpoint.
   INFO: ClassVar[ModelInfo] = ModelInfo()
 
   def __post_init__(self):
@@ -228,7 +228,7 @@ def vision_encoder(self) -> gemma4_vision.VisionEncoder | None:
           'audio_soft_token_counts',
       ),
   )
-  # The function accepts/returns aribtrary batch shape, but inside the
+  # The function accepts/returns arbitrary batch shape, but inside the
   # function, the batch dimension is flattened to a single dimension.
   @_jax_utils.flatten_unflatten_batch_dim()
   @typechecked

gemma/gm/nn/gemma4/audio/_modules.py

@@ -403,8 +403,8 @@ class LocalDotProductAttention(nn.Module):
   the SequenceLayers library. The latter is made for streaming attention and
   requires a specific kind of tensor manipulation. In order to match exactly
   output of the internal implementation of the AudioTokenizer,
-  we perform those manipulation here even though they are not theoritically
-  necessary. This results in a complexified implementation of a theoritically
+  we perform those manipulation here even though they are not theoretically
+  necessary. This results in a complexified implementation of a theoretically
   simple attention mechanism.
   """
 

gemma/gm/nn/vision/_vision.py

@@ -50,7 +50,7 @@ def check_mask(input_data: Float["L"]) -> tuple[Bool["1"], Int["L"]]:
 
   Returns:
     A boolean indicating whether the mask contains the correct number of blocks
-    and there starting positions (filled with 0 for jit compatiblity).
+    and there starting positions (filled with 0 for jit compatibility).
   """
   is_mask = input_data == TOKEN_PLACEHOLDER
   start_idx = (
@@ -124,7 +124,7 @@ def initialize_vision_tokens(
 ) -> VisionInitEmbeddings:
   """Initializes vision embeddings.
 
-  Vision data intialization wrapper for sampling.
+  Vision data initialization wrapper for sampling.
 
   Example (text only inference):
   ```python
@@ -191,7 +191,7 @@ def initialize_vision_tokens(
     base_mm_tokens = base_mm_tokens.at[:, -1].set(NEW_LINE_TOKEN)
     # then we repeat this tensor for each image in the batch
     mm_tokens = jnp.repeat(base_mm_tokens, patches.shape[1], axis=0)
-    # insert iamges after the first token, which is BOS as expected in gemma v3
+    # insert images after the first token, which is BOS as expected in gemma v3
     token_buffer = jnp.concatenate(
         [token_buffer[:, :1], mm_tokens, token_buffer[:, 1:]], axis=1
     )

gemma/gm/text/_sampler.py

@@ -118,7 +118,7 @@ class Sampler:
       conversation can have (prompts, answers, images for all turns). Setting
       this to a fixed value avoids re-compilation between turns.
     max_out_length: Length of the output buffer for a single turn. Static value
-      used to avoid trigering a jit recompilation. Shouldn't be changed unless
+      used to avoid triggering a jit recompilation. Shouldn't be changed unless
       you have a task where the model generates really long outputs.
     pad_length: If provided, pad the prompt to this length. This ensure the
       prompt is always the same length, to avoid jit re-compilation.

gemma/gm/text/_tokenizer.py

@@ -41,7 +41,7 @@
   import plotly.express as px  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
 
 
-_WHITESPACE_CHAR = '▁'  # Note this is NOT a undescore (▁ != _)
+_WHITESPACE_CHAR = '▁'  # Note this is NOT a underscore (▁ != _)
 
 
 class _DisplayEnumType(enum.EnumType):

gemma/gm/text/_tool_sampler.py

@@ -90,7 +90,7 @@ def chat(
         Overrites the `multi_turn` attribute.
       is_legacy_tool_answer: When `True`, indicates that the model has emitted
         `<eos>` rather than `<|tool_response>`, thus this needs to be corrected.
-        (this is an internal variable that should never be explictly set).
+        (this is an internal variable that should never be explicitly set).
       sharding: Sharding tree (Gemma 4 only).
 
     Returns:

gemma/peft/_einsum_utils.py

@@ -28,7 +28,7 @@ def get_lora_einsum_str_and_shapes(
 ) -> tuple[str, _Shape, _Shape]:
   """Extract the LoRA decomposition from the original einsum parameters.
 
-  This function reqrites a einsum string `inputs,weights->outputs` into
+  This function rewrites a einsum string `inputs,weights->outputs` into
   `inputs,a,b->outputs`.
 
   Args:
@@ -66,7 +66,7 @@ def get_lora_einsum_str_and_shapes(
 
   lora_einsum_str = f'{inputs},{a_str},{b_str}->{outputs}'
 
-  # This assume there's no elipsis in the weights.
+  # This assume there's no ellipsis in the weights.
   weights_str_to_dim = dict(zip(weights, weights_shape))
   weights_str_to_dim[rank_dim] = rank
   a_shape = tuple(weights_str_to_dim[c] for c in a_str)

gemma/peft/_einsum_utils_test.py

@@ -40,7 +40,7 @@ def test_einsum():
       rank=R,
   ) == ('bijk,bjkr,rnm->bimn', (B, J, K, R), (R, N, M))
 
-  # When r is already in use in the letters, another letter is choosen
+  # When r is already in use in the letters, another letter is chosen
   # (here `D`)
   assert _einsum_utils.get_lora_einsum_str_and_shapes(
       einsum_str='...abcABCrR,Rgk->...ABCabcrgk',

gemma/peft/_quantization_utils.py

@@ -71,7 +71,7 @@ def quantize(
 ) -> PyTree:
   """Quantizes the given params.
 
-  In ths API, we convert the elements of params in order to actually get
+  In the API, we convert the elements of params in order to actually get
   quantized values. It is currently limited to INT$ per-channel weight
   quantization.
 
@@ -146,7 +146,7 @@ def convert_leaf(data: Any) -> Any:
       data = quantize_leaf(data, checkpoint_kernel_key)
     # This hack is required because the FeedForward layer call two different
     # Einsum with using `nn.share_scope`, so the two wrappers need a different
-    # name. Weights are not stored under any kernel key and are isntead under
+    # name. Weights are not stored under any kernel key and are instead under
     # the following names.
     if 'gating_einsum' in data or 'linear' in data:
       data = quantize_leaf(data, 'gating_einsum')

gemma/research/t5gemma/README.md

@@ -2,7 +2,7 @@
 
 T5Gemma (aka encoder-decoder Gemma) was proposed in a
 [research paper](https://arxiv.org/abs/2504.06225) by Google. It is a family of
-encoder-decoder large langauge models, developed by adapting pretrained
+encoder-decoder large language models, developed by adapting pretrained
 decoder-only models into encoder-decoder. T5Gemma includes pretrained and
 instruction-tuned variants. The architecture is based on transformer
 encoder-decoder design following T5, with improvements from Gemma 2: GQA, RoPE,

pyproject.toml

@@ -67,7 +67,7 @@ changelog = "https://github.com/google-deepmind/gemma/blob/main/CHANGELOG.md"
 # documentation = ""
 
 [project.optional-dependencies]
-# Development deps (unittest, linting, formating,...)
+# Development deps (unittest, linting, formatting,...)
 # Installed through `pip install -e .[dev]`
 dev = [
     "pytest",