Ignore the serialized Mosaic kernel when diffing HLO

The Mosaic module of a Pallas custom call is serialized into the HLO and carries the absolute
path of the installed jax package, so the dump differs between a venv install and a Docker image
install even though the compiled graph is identical. That made hlo_diff_test fail on the Docker
image CI jobs. Normalize the payload away in filter_line, like stack_frame_id already is.

Author: ecnal-cienet

tests/integration/hlo_diff_test.py

@@ -43,13 +43,21 @@
 # Filters and normalizes an HLO line to ignore environment-specific details.
 STACK_FRAME_REGEX = re.compile(r"stack_frame_id=\d+")
 
+# The serialized Mosaic kernel of a Pallas custom call embeds the absolute path of the installed
+# jax package, so it differs between a venv install and a Docker image install even though the
+# compiled graph is the same. The kernel's shapes and block sizes are still checked, since they
+# appear as plain text elsewhere in the dump.
+CUSTOM_CALL_BODY_REGEX = re.compile(r'("custom_call_config":\{"body":")[^"]*(")')
+
 
 def filter_line(line):
   # Matches operation sections content like: 1234 {"sharding parameters", ...} or 1234 "operation metadata"
   if SECTION_CONTENT_REGEX.match(line):
     return None
   # Replace stack_frame_id=... with stack_frame_id=0
   line = STACK_FRAME_REGEX.sub("stack_frame_id=0", line)
+  # Replace the serialized Mosaic kernel with a placeholder
+  line = CUSTOM_CALL_BODY_REGEX.sub(r"\1<mosaic_body>\2", line)
   return line