diff --git a/gemma/colabs/fine_tuning_tutorial.ipynb b/colabs/fine_tuning_tutorial.ipynb
similarity index 100%
rename from gemma/colabs/fine_tuning_tutorial.ipynb
rename to colabs/fine_tuning_tutorial.ipynb
diff --git a/gemma/colabs/gsm8k_eval.ipynb b/colabs/gsm8k_eval.ipynb
similarity index 100%
rename from gemma/colabs/gsm8k_eval.ipynb
rename to colabs/gsm8k_eval.ipynb
diff --git a/gemma/colabs/sampling_tutorial.ipynb b/colabs/sampling_tutorial.ipynb
similarity index 100%
rename from gemma/colabs/sampling_tutorial.ipynb
rename to colabs/sampling_tutorial.ipynb
diff --git a/gemma/examples/sampling.py b/examples/sampling.py
similarity index 81%
rename from gemma/examples/sampling.py
rename to examples/sampling.py
index 5dd4f9c..277228c 100644
--- a/gemma/examples/sampling.py
+++ b/examples/sampling.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 r"""An example showing how to load a checkpoint and sample from it.
 
 Getting Started with Gemma Sampling:
@@ -42,8 +28,9 @@
 
 python sampling.py --path_checkpoint=${PATH_TO_THE_GEMMA_CHECKPOINT} \
     --path_tokenizer=${PATH_TO_THE_GEMMA_TOKENIZER} \
-    --string_to_sample="Where is Paris?" 
+    --string_to_sample="Where is Paris?"
 """
+
 from typing import Sequence
 
 from absl import app
@@ -91,7 +78,7 @@ def _load_and_sample(
   vocab.Load(path_tokenizer)
   transformer_config = transformer_lib.TransformerConfig.from_params(
       parameters,
-      cache_size=cache_size
+      cache_size=cache_size,
   )
   transformer = transformer_lib.Transformer(transformer_config)
   sampler = sampler_lib.Sampler(
diff --git a/gemma/__init__.py b/gemma/__init__.py
index a58f9e8..ba61169 100644
--- a/gemma/__init__.py
+++ b/gemma/__init__.py
@@ -12,21 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
-
 """gemma API."""
 
 # A new PyPI release will be pushed every time `__version__` is increased.
diff --git a/gemma/layers.py b/gemma/layers.py
index 6fd4077..f6179cb 100644
--- a/gemma/layers.py
+++ b/gemma/layers.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Base layers."""
 
 from flax import linen as nn
diff --git a/gemma/layers_test.py b/gemma/layers_test.py
index 7770633..2c50c32 100644
--- a/gemma/layers_test.py
+++ b/gemma/layers_test.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Tests for transformer layers."""
 
 from absl.testing import absltest
diff --git a/gemma/modules.py b/gemma/modules.py
index 392cc41..7a1fa67 100644
--- a/gemma/modules.py
+++ b/gemma/modules.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Transformer sub-modules."""
 
 import enum
diff --git a/gemma/modules_test.py b/gemma/modules_test.py
index 216cded..573bd2a 100644
--- a/gemma/modules_test.py
+++ b/gemma/modules_test.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Tests for transformer modules."""
 
 import logging
diff --git a/gemma/params.py b/gemma/params.py
index 657b91e..1701632 100644
--- a/gemma/params.py
+++ b/gemma/params.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Utils for loading Gemma params."""
 
 import functools
diff --git a/gemma/params_test.py b/gemma/params_test.py
index 8e1c170..74b9f48 100644
--- a/gemma/params_test.py
+++ b/gemma/params_test.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Tests for transformer params."""
 
 from absl.testing import absltest
diff --git a/gemma/positional_embeddings.py b/gemma/positional_embeddings.py
index 2bcab16..6d3527c 100644
--- a/gemma/positional_embeddings.py
+++ b/gemma/positional_embeddings.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Utils for positional embeddings (including RoPE)."""
 
 import jax
diff --git a/gemma/positional_embeddings_test.py b/gemma/positional_embeddings_test.py
index 24c8f04..f039fa1 100644
--- a/gemma/positional_embeddings_test.py
+++ b/gemma/positional_embeddings_test.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Tests for the positional embeddings utilities."""
 
 from absl.testing import absltest
diff --git a/gemma/sampler.py b/gemma/sampler.py
index 3786cc7..1ec029c 100644
--- a/gemma/sampler.py
+++ b/gemma/sampler.py
@@ -12,24 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Sampler for Gemma transformer.
 
 An example of a sampling class for a Gemma model.
 """
+
 from collections.abc import Sequence
 import dataclasses
 
diff --git a/gemma/sampler_test.py b/gemma/sampler_test.py
index 8c5edd0..cadaef5 100644
--- a/gemma/sampler_test.py
+++ b/gemma/sampler_test.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Minimal test for sampler."""
 
 from typing import Iterable
diff --git a/gemma/transformer.py b/gemma/transformer.py
index 52e34ad..f75bf9f 100644
--- a/gemma/transformer.py
+++ b/gemma/transformer.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Gemma transformer."""
 
 import dataclasses
diff --git a/gemma/transformer_test.py b/gemma/transformer_test.py
index 2284ff2..6f34a9e 100644
--- a/gemma/transformer_test.py
+++ b/gemma/transformer_test.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2024 DeepMind Technologies Limited.
-#
-# 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.
-# ============================================================================
 """Tests for the Gemma transformer."""
 
 from unittest import mock