Merge pull request #28899 from jenriver:absl_logging_fix

PiperOrigin-RevId: 762099728

Author: Google-ML-Automation

docs/export/export.md

@@ -710,10 +710,7 @@ total 32
 -rw-rw-r--@ 1 necula  wheel  2333 Jun 19 11:04 jax_ir3_jit_my_fun_export.mlir
 ```
 
-Inside Google, you can turn on logging by using the `--vmodule` argument to
-specify the logging levels for different modules,
-e.g., `--vmodule=_export=3`.
-
+Set [`JAX_DEBUG_LOG_MODULES=jax._src.export`](https://docs.jax.dev/en/latest/config_options.html#jax_debug_log_modules) to enable extra debugging logging.
 
 (export_ensuring_compat)=
 ### Ensuring forward and backward compatibility

jax/_src/compiler.py

@@ -241,7 +241,7 @@ def get_compile_options(
   else:
     compile_options.profile_version = _NO_PROFILE_DONT_RETRIEVE
     if backend is None:
-      logging.info("get_compile_options: no backend supplied; "
+      logger.info("get_compile_options: no backend supplied; "
                    "disabling XLA-AutoFDO profile")
     else:
       fdo_profile_version = get_latest_profile_version(backend)
@@ -369,7 +369,7 @@ def compile_or_get_cached(
   module_name = ir.StringAttr(sym_name).value
 
   if dumped_to := mlir.dump_module_to_file(computation, "compile"):
-    logging.info("Dumped the module to %s.", dumped_to)
+    logger.info("Dumped the module to %s.", dumped_to)
 
   is_multi_process = (
       len({device.process_index for device in devices.flatten()}) > 1
@@ -514,7 +514,7 @@ def _resolve_compilation_strategy(
     # The compilation cache is enabled and AutoPGLE is enabled/expected
     if _is_executable_in_cache(backend, pgle_optimized_cache_key):
       if config.compilation_cache_expect_pgle.value:
-        logging.info(f"PGLE-optimized {module_name} loaded from compilation cache")
+        logger.info(f"PGLE-optimized {module_name} loaded from compilation cache")
       # No need to record N profiles in this case
       if pgle_profiler is not None:
         pgle_profiler.disable()

jax/_src/export/_export.py

@@ -25,7 +25,7 @@
 import re
 from typing import Any, Protocol, TypeVar, Union, cast
 
-from absl import logging
+import logging
 import numpy as np
 
 import jax
@@ -55,6 +55,8 @@
 
 from jax._src.export import shape_poly
 
+logger = logging.getLogger(__name__)
+
 map = util.safe_map
 zip = util.safe_zip
 
@@ -704,16 +706,15 @@ def _export_lowered(
     out_avals_flat = lowered.compile_args["out_avals"]  # type: ignore
 
   # Log and then check the module.
-  if logging.vlog_is_on(3):
-    logmsg = (f"fun_name={fun_name} version={version} "
-              f"lowering_platforms={lowering._platforms} "  # type: ignore[unused-ignore,attribute-error]
-              f"disabled_checks={disabled_checks}")
-    logging.info("Exported JAX function: %s\n", logmsg)
-    logging.info(mlir.dump_module_message(mlir_module, "export"))
-    logging.info(
-        "Size of mlir_module_serialized: %d byte",
-        len(mlir_module_serialized),
-    )
+  logmsg = (f"fun_name={fun_name} version={version} "
+            f"lowering_platforms={lowering._platforms} "  # type: ignore[unused-ignore,attribute-error]
+            f"disabled_checks={disabled_checks}")
+  logger.debug("Exported JAX function: %s\n", logmsg)
+  logger.debug(mlir.dump_module_message(mlir_module, "export"))
+  logger.debug(
+      "Size of mlir_module_serialized: %d byte",
+      len(mlir_module_serialized),
+  )
 
   _check_module(mlir_module,
                 disabled_checks=disabled_checks,