Leverage IR/HLO save tensor utils for Neuron

pytorch · Jan 11, 2025 · 2bf7575 · 2bf7575
1 parent 6676b37
commit 2bf7575
Showing 1 changed file with 3 additions and 13 deletions.
diff --git a/test/neuron/run_tests.sh b/test/neuron/run_tests.sh
@@ -96,16 +96,6 @@ function run_eager_debug {
   XLA_USE_EAGER_DEBUG_MODE=1 run_test "$@"
 }
 
-function run_save_tensor_ir {
-  echo "Running in save tensor file mode: $@"
-  XLA_SAVE_TENSORS_FILE="/tmp/xla_test_save_ir.txt" XLA_SAVE_TENSORS_FMT="text" run_test "$@"
-}
-
-function run_save_tensor_hlo {
-  echo "Running in save tensor file mode: $@"
-  XLA_SAVE_TENSORS_FILE="/tmp/xla_test_save_ir.txt" XLA_SAVE_TENSORS_FMT="hlo" run_test "$@"
-}
-
 function run_pt_xla_debug {
   echo "Running in save tensor file mode: $@"
   PT_XLA_DEBUG=1 PT_XLA_DEBUG_FILE="/tmp/pt_xla_debug.txt" run_test "$@"
@@ -169,16 +159,16 @@ function run_xla_op_tests1 {
   run_test "$CDIR/dynamo/test_num_output.py"
   run_test "$CDIR/dynamo/test_graph_input_matcher.py"
   run_test "$CDIR/dynamo/test_dynamo_config.py"
-  run_save_tensor_ir "$CDIR/dynamo/test_dynamo_graph_dump.py"
+  run_save_tensor_ir run_test "$CDIR/dynamo/test_dynamo_graph_dump.py"
   #run_test "$CDIR/test_data_type.py"
   run_use_bf16 "$CDIR/test_data_type.py"
   run_downcast_bf16 "$CDIR/test_data_type.py"
   #run_test "$CDIR/test_fp8.py"
   run_xla_ir_debug "$CDIR/test_env_var_mapper.py"
   run_xla_hlo_debug "$CDIR/test_env_var_mapper.py"
   run_xla_hlo_debug "$CDIR/stablehlo/test_stablehlo_save_load.py"
-  run_save_tensor_ir "$CDIR/spmd/test_spmd_graph_dump.py"
-  run_save_tensor_hlo "$CDIR/spmd/test_spmd_graph_dump.py"
+  run_save_tensor_ir run_test "$CDIR/spmd/test_spmd_graph_dump.py"
+  run_save_tensor_hlo run_test "$CDIR/spmd/test_spmd_graph_dump.py"
 }
 
 function run_xla_op_tests2 {