Formalize the changes.

Author: RoshaniN

.github/workflows/RunTests.yml

@@ -75,9 +75,6 @@ jobs:
       device_name: v4-8
       pytest_marker: 'not gpu_only and not integration_test'
       test_directory: 'tests'
-      # xla_python_client_mem_fraction: 0.75
-      # tf_force_gpu_allow_growth: false
-      # container_resource_option: "--privileged"
 
   tpu_integration_tests:
     needs: tpu_image

MaxText/tests/train_tests.py

@@ -166,8 +166,6 @@ def test_tpu_dropout(self):
   def test_gpu_dropout(self):
     train_main(TrainTests.CONFIGS["dropout"] + ["attention=dot_product"])
 
-  # # TODO (b/393393501) : MaxText build failure in hf_data_processing_test.py
-  # @pytest.mark.skip(reason="Tests are currently flaking / failing due to HF token issues")
   @pytest.mark.tpu_only
   def test_tpu_hf_input_pipeline(self):
     train_main(TrainTests.CONFIGS["hf_input_pipeline"])

docker_run_pathways_containers.sh

@@ -14,18 +14,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Example - bash docker_run_pathways_containers.sh maxtext_image=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest command="cd MaxText ; python3 -m pytest tests -m 'not gpu_only and not integration_test' -s"
-
-# Stop execution if any command exits with error
-
+# Examples -
+# TPU unit tests - bash docker_run_pathways_containers.sh maxtext_image=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest command="cd MaxText ; python3 -m pytest tests -m 'not gpu_only and not integration_test' -s"
 
 echo "Running docker_run_pathways_containers.sh"
 
+# Stop execution if any command exits with error
 set -e
 
-# Defaults -
-maxtext_image=us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
-command="cd MaxText ; python3 -m pytest tests -m 'not gpu_only and not integration_test' -s"
 
 # Parse input variables
 for ARGUMENT in "$@"; do