A call to logging.info for each line of aot_model_analysis result

GitOrigin-RevId: 8d4fbec

Author: wangkuiyi

axlearn/common/trainer.py

@@ -1231,7 +1231,8 @@ def compile_train_step(
             # pjit currently requires all parameters to be specified as positional args.
             lowered_train_step = jit_train_step.lower(trainer_state, input_batch)
             compiled = lowered_train_step.compile(compiler_options=compiler_options)
-            logging.log_first_n(logging.INFO, aot_model_analysis(compiled), 1)
+            for line in aot_model_analysis(compiled).split("\n"):
+                logging.info("%s", line)
             return compiled
 
     def _train_step(