Merge pull request #2873 from AI-Hypercomputer:nicogrande/improve-max-logging

PiperOrigin-RevId: 848356647

Author: Google-ML-Automation

src/MaxText/max_logging.py

@@ -12,8 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Stub for logging utilities. Right now just meant to avoid raw prints."""
+"""Logging utilities."""
+from absl import logging
 
 
 def log(user_str):
-  print(user_str, flush=True)
+  """Logs a message at the INFO level."""
+  # Note, stacklevel=2 makes the log show the caller of this function.
+  logging.info(user_str, stacklevel=2)
+
+
+def debug(user_str):
+  """Logs a message at the DEBUG level."""
+  logging.debug(user_str, stacklevel=2)
+
+
+def info(user_str):
+  """Logs a message at the INFO level."""
+  logging.info(user_str, stacklevel=2)
+
+
+def warning(user_str):
+  """Logs a message at the WARNING level."""
+  logging.warning(user_str, stacklevel=2)
+
+
+def error(user_str):
+  """Logs a message at the ERROR level."""
+  logging.error(user_str, stacklevel=2)