Adjust copybara rules

PiperOrigin-RevId: 864637021

Author: Google-ML-Automation

tests/integration/smoke/train_gpu_smoke_test.py

@@ -21,7 +21,7 @@
 from maxtext.common.gcloud_stub import is_decoupled
 from MaxText.globals import MAXTEXT_ASSETS_ROOT, MAXTEXT_PKG_DIR
 from MaxText.train import main as train_main
-from maxtext.tests.utils.test_helpers import get_test_dataset_path, get_test_base_output_directory
+from tests.utils.test_helpers import get_test_dataset_path, get_test_base_output_directory
 
 
 class Train(unittest.TestCase):

tests/integration/smoke/train_int8_smoke_test.py

@@ -21,7 +21,7 @@
 from maxtext.common.gcloud_stub import is_decoupled
 from MaxText.globals import MAXTEXT_ASSETS_ROOT
 from MaxText.train import main as train_main
-from maxtext.tests.utils.test_helpers import get_test_config_path, get_test_dataset_path, get_test_base_output_directory
+from tests.utils.test_helpers import get_test_config_path, get_test_dataset_path, get_test_base_output_directory
 
 
 class Train(unittest.TestCase):

tests/integration/smoke/train_smoke_test.py

@@ -18,7 +18,7 @@
 
 from absl.testing import absltest
 
-from maxtext.tests.utils.test_helpers import get_test_config_path, get_test_dataset_path, get_test_base_output_directory
+from tests.utils.test_helpers import get_test_config_path, get_test_dataset_path, get_test_base_output_directory
 from MaxText.train import main as train_main
 from MaxText.globals import MAXTEXT_PKG_DIR, MAXTEXT_ASSETS_ROOT
 from maxtext.common.gcloud_stub import is_decoupled

tests/integration/smoke/train_using_ragged_dot_smoke_test.py

@@ -20,7 +20,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 
-from maxtext.tests.utils.test_helpers import get_test_config_path
+from tests.utils.test_helpers import get_test_config_path
 from MaxText import globals as maxtext_globals, train
 
 train_main = train.main

tests/utils/attention_test_util.py

@@ -30,7 +30,7 @@
 from MaxText.globals import MAXTEXT_PKG_DIR
 from MaxText.layers.attention_mla import MLA
 from MaxText.sharding import maybe_shard_with_name
-from maxtext.tests.utils.test_helpers import get_test_config_path
+from tests.utils.test_helpers import get_test_config_path
 
 
 class MLATestBase(parameterized.TestCase):