diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ab8542500bf..51a879f1a662 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,9 @@ Remember to align the itemized text with the first line of an item within a list ## jaxlib 0.4.23 +* Fixed a bug that caused verbose logging from the GPU compiler during + compilation. + ## jax 0.4.22 (Dec 13, 2023) * Deprecations diff --git a/third_party/xla/log.patch b/third_party/xla/log.patch new file mode 100644 index 000000000000..959e1cd53c42 --- /dev/null +++ b/third_party/xla/log.patch @@ -0,0 +1,13 @@ +diff --git a/xla/service/gpu/gpu_compiler.cc b/xla/service/gpu/gpu_compiler.cc +index 302879c0f..9c55820c5 100644 +--- a/xla/service/gpu/gpu_compiler.cc ++++ b/xla/service/gpu/gpu_compiler.cc +@@ -546,7 +546,7 @@ Status GpuCompiler::OptimizeHloModule(HloModule* hlo_module, + + // LOG_LINES is used instead of LOG since the message can exceed the + // maximum line length, which results in the message being truncated. +- XLA_LOG_LINES( ++ XLA_VLOG_LINES( + 1, absl::StrFormat("GpuCompilationEnvironment of hlo_module %s:\n%s", + hlo_module->name(), debug_options.DebugString())); + diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl index 3c073398093a..a6f972c34943 100644 --- a/third_party/xla/workspace.bzl +++ b/third_party/xla/workspace.bzl @@ -29,6 +29,7 @@ def repo(): sha256 = XLA_SHA256, strip_prefix = "xla-{commit}".format(commit = XLA_COMMIT), urls = tf_mirror_urls("https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)), + patch_file = ["//third_party/xla:log.patch"], ) # For development, one often wants to make changes to the TF repository as well