Merge pull request #2956 from AI-Hypercomputer:distillation_restructure

PiperOrigin-RevId: 859237519

Author: Google-ML-Automation

.coveragerc

@@ -9,8 +9,10 @@ omit =
 
 [paths]
 source =
+    src/MaxText
     src/MaxText
     */site-packages/MaxText
+    */site-packages/maxtext
 
 [report]
 show_missing = True

.github/workflows/run_pathways_tests.yml

@@ -97,7 +97,8 @@ jobs:
           export MAXTEXT_PKG_DIR=$(pwd)/src/MaxText
           # TODO(b/454659463): Enable test_default_hlo_match after volume mount is supported.
           .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not AotHloIdenticalTest and not CompileThenLoad" --durations=0
-    
+        env:
+          PYTHONPATH: "${{ github.workspace }}/src"
     services:
       resource_manager:
         image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest

.github/workflows/run_tests_against_package.yml

@@ -122,6 +122,7 @@ jobs:
             --durations=0 \
             --deselect "tests/unit/tokenizer_test.py::TokenizerTest::test_detokenize" \
             --cov=MaxText \
+            --cov=maxtext \
             --cov-report=xml \
             --cov-report=term \
             $SPLIT_ARGS

codecov.yml

@@ -39,6 +39,7 @@ ignore:
   - "src/MaxText/inference"
   - "src/MaxText/inference_mlperf"
   - "src/MaxText/scratch_code"
+  - "src/MaxText/distillation" # code moved to src/MaxText/trainers/post_train/distillation
 
 
 flags:

src/MaxText/distillation/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023–2025 Google LLC
+# Copyright 2023–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.