Export extras directory to repository root on GitHub.

pathwaysutils/extras/elastic/README.md → extras/elastic/README.md

@@ -4,7 +4,7 @@ Evaluating distributed machine learning training on elastic TPU topologies requi
 
 ## Live Interactive Simulation
 
-* [Hosted GitHub Pages Simulation](https://ai-hypercomputer.github.io/pathways-utils/pathwaysutils/extras/elastic/model_comparison.html) is a live interactive web application with simulation controls, or [HTMLPreview Instant Mirror](https://htmlpreview.github.io/?https://github.com/AI-Hypercomputer/pathways-utils/blob/main/pathwaysutils/extras/elastic/model_comparison.html).
+* [Hosted GitHub Pages Simulation](https://ai-hypercomputer.github.io/pathways-utils/extras/elastic/model_comparison.html) is a live interactive web application with simulation controls, or [HTMLPreview Instant Mirror](https://htmlpreview.github.io/?https://github.com/AI-Hypercomputer/pathways-utils/blob/main/extras/elastic/model_comparison.html).
 * Local browser viewing allows running `model_comparison.html` directly on your workstation by opening the file in Chrome, Firefox, or Safari, requiring zero installation or server dependencies.
 
 ## Recovery Paradigms Compared

extras/simple_model/__init__.py

@@ -0,0 +1,48 @@
+# Copyright 2026 Google LLC
+#
+# 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
+#
+#      https://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.
+"""Simple Transformer model package using Flax NNX and Grain for Pathways."""
+
+from extras.simple_model.compile import CompiledStepArtifacts
+from extras.simple_model.compile import compile_training_step
+from extras.simple_model.compile import compute_mesh_shape
+from extras.simple_model.compile import update_step
+from extras.simple_model.config import CONFIGS
+from extras.simple_model.config import DEFAULT_CONFIG
+from extras.simple_model.config import ModelConfig
+from extras.simple_model.config import parse_topology
+from extras.simple_model.data import create_data_iterator
+from extras.simple_model.data import create_token_dataset
+from extras.simple_model.data import SyntheticTokenSource
+from extras.simple_model.model import geglu
+from extras.simple_model.model import SimpleModel
+from extras.simple_model.model import TransformerBlock
+from extras.simple_model.workload import Workload
+
+__all__ = [
+    "CONFIGS",
+    "CompiledStepArtifacts",
+    "DEFAULT_CONFIG",
+    "ModelConfig",
+    "SimpleModel",
+    "SyntheticTokenSource",
+    "TransformerBlock",
+    "Workload",
+    "compile_training_step",
+    "compute_mesh_shape",
+    "create_data_iterator",
+    "create_token_dataset",
+    "geglu",
+    "parse_topology",
+    "update_step",
+]

pathwaysutils/extras/simple_model/compile.py → extras/simple_model/compile.py

@@ -23,8 +23,8 @@
 import jax
 from jax import numpy as jnp
 import numpy as np
-from pathwaysutils.extras.simple_model import config as config_lib
-from pathwaysutils.extras.simple_model import model as model_lib
+from extras.simple_model import config as config_lib
+from extras.simple_model import model as model_lib
 
 
 def compute_mesh_shape(

pathwaysutils/extras/simple_model/data.py → extras/simple_model/data.py

@@ -17,7 +17,7 @@
 
 import grain.python as grain
 import numpy as np
-from pathwaysutils.extras.simple_model import config as config_lib
+from extras.simple_model import config as config_lib
 
 
 class SyntheticTokenSource(grain.RandomAccessDataSource):

pathwaysutils/extras/simple_model/model.py → extras/simple_model/model.py

@@ -16,7 +16,7 @@
 from flax import nnx
 import jax
 from jax import numpy as jnp
-from pathwaysutils.extras.simple_model import config as config_lib
+from extras.simple_model import config as config_lib
 
 
 def geglu(x: jax.Array) -> jax.Array:

pathwaysutils/extras/simple_model/simple_model.py → extras/simple_model/simple_model.py

@@ -19,9 +19,9 @@
 from absl import app
 from absl import flags
 import pathwaysutils
-from pathwaysutils.extras.simple_model.config import DEFAULT_CONFIG
-from pathwaysutils.extras.simple_model.config import parse_topology
-from pathwaysutils.extras.simple_model.workload import Workload
+from extras.simple_model.config import DEFAULT_CONFIG
+from extras.simple_model.config import parse_topology
+from extras.simple_model.workload import Workload
 
 _TOPOLOGY = flags.DEFINE_string(
     "topology", "", "Optional TPU slice topology (e.g., 4x8, 4x4x4)."

pathwaysutils/extras/simple_model/simple_model_test.py → extras/simple_model/simple_model_test.py

@@ -23,11 +23,11 @@
 from jax import numpy as jnp
 import numpy as np
 from pathwaysutils.elastic import manager
-from pathwaysutils.extras.simple_model import compile as compile_lib
-from pathwaysutils.extras.simple_model import config as config_lib
-from pathwaysutils.extras.simple_model import data as data_lib
-from pathwaysutils.extras.simple_model import model as model_lib
-from pathwaysutils.extras.simple_model import workload as workload_lib
+from extras.simple_model import compile as compile_lib
+from extras.simple_model import config as config_lib
+from extras.simple_model import data as data_lib
+from extras.simple_model import model as model_lib
+from extras.simple_model import workload as workload_lib
 
 
 def setUpModule():

pathwaysutils/extras/simple_model/workload.py → extras/simple_model/workload.py

@@ -24,10 +24,10 @@
 from jax import numpy as jnp
 import orbax.checkpoint as ocp
 from pathwaysutils.elastic import manager
-from pathwaysutils.extras.simple_model import compile as compile_lib
-from pathwaysutils.extras.simple_model import config as config_lib
-from pathwaysutils.extras.simple_model import data as data_lib
-from pathwaysutils.extras.simple_model import model as model_lib
+from extras.simple_model import compile as compile_lib
+from extras.simple_model import config as config_lib
+from extras.simple_model import data as data_lib
+from extras.simple_model import model as model_lib
 
 _logger = logging.getLogger(__name__)