Repository: google/jax
Pull request #23421 (2024-09-03 22:28:12)

Add frontend attributes to Jax. This allows Jax users to annotate Jax code with frontend_attributes which can be traced down to the HLO level, to be used for numerical debugging purposes.
Add frontend attributes to Jax. This allows Jax users to annotate Jax code with frontend_attributes which can be traced down to the HLO level, to be used for numerical debugging purposes.



Patch:
diff --git a/jax/BUILD b/jax/BUILD
index 4c6221949..74072dc44 100644
--- a/jax/BUILD
+++ b/jax/BUILD
@@ -299,6 +299,7 @@ py_library_providing_imports_info(
         ":version",
         ":xla",
         ":xla_bridge",
+        ":xla_metadata",
         "//jax/_src/lib",
     ] + py_deps("numpy") + py_deps("scipy") + py_deps("opt_einsum") + py_deps("flatbuffers") + jax_extra_deps,
 )
@@ -451,6 +452,7 @@ pytype_strict_library(
         ":tree_util",
         ":typing",
         ":util",
+        ":xla_metadata",
         "//jax/_src/lib",
     ] + py_deps("numpy"),
 )
@@ -703,6 +705,7 @@ pytype_strict_library(
         ":state_types",
         ":tree_util",
         ":util",
+        ":xla_metadata",
     ] + py_deps("numpy"),
 )
 
@@ -768,6 +771,14 @@ pytype_strict_library(
     deps = [":config"],
 )
 
+pytype_strict_library(
+    name = "xla_metadata",
+    srcs = ["_src/xla_metadata.py"],
+    deps = [
+        ":config",
+    ],
+)
+
 pytype_strict_library(
     name = "layout",
     srcs = ["_src/layout.py"],
diff --git a/jax/_src/config.py b/jax/_src/config.py
index 646e487c5..b2d1aa52e 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -202,16 +202,20 @@ def trace_context():
   tls = jax_jit.thread_local_state()
   axis_env_state = ()
   mesh_context_manager = ()
+  xla_metadata_context_manager = ()
   compute_on_context_manager = ()
+
   context: Any = tls.extra_jit_context
   if context and context.axis_env_state is not None:
     axis_env_state = context.axis_env_state
   if context and context.mesh_context_manager:
     mesh_context_manager = context.mesh_context_manager
+  if context and context.xla_metadata_context_manager:
+    xla_metadata_context_manager = context.xla_metadata_context_manager
   if context and context.compute_on_context_manager:
     compute_on_context_manager = context.compute_on_context_manager
-  return (axis_env_state, mesh_context_manager, compute_on_context_manager,
-          enable_x64.value,
+  return (axis_env_state, mesh_context_manager, xla_metadata_context_manager,
+          compute_on_context_manager, enable_x64.value,
           numpy_rank_promotion.value, default_matmul_precision.value,
           dynamic_shapes.value, numpy_dtype_promotion.value,
           default_device.value, random_seed_offset.value,
@@ -858,6 +862,7 @@ class _ThreadLocalExtraJitContext(NamedTuple):
   axis_env_state: Hashable = ()
   mesh_context_manager: Hashable = ()
   compute_on_context_manager: Hashable = ()
+  xla_metadata_context_manager: Hashable = ()
 
   # Values set by _StateContextManager context managers.
   # CAUTION: these must be initialized to `None`! The state context manager
diff --git a/jax/_src/core.py b/jax/_src/core.py
index f6da7ac7a..ef3ace2e0 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -55,6 +55,7 @@ from jax._src.lib import jax_jit
 from jax._src import traceback_util
 from jax._src.typing import Array, DimSize, Shape
 from jax._src import typing
+from jax._src import xla_metadata as xla_metadata_lib
 
 traceback_util.register_exclusion(__file__)
 
@@ -261,12 +262,15 @@ def jaxpr_as_fun(closed_jaxpr: ClosedJaxpr, *args):
 
 class JaxprEqnContext:
 
-  def __init__(self, compute_type: str | None, threefry_partitionable: bool):
+  def __init__(self, compute_type: str | None, threefry_partitionable: bool,
+               xla_metadata=None):
     self.compute_type = compute_type
     self.threefry_partitionable = threefry_partitionable
+    self.xla_metadata = xla_metadata
     self._managers = [
         (compute_on.extend_compute_type, self.compute_type),
         (config.threefry_partitionable.__call__, self.threefry_partitionable),
+        (xla_metadata_lib.set_xla_metadata, self.xla_metadata),
     ]
 
   @property
@@ -278,8 +282,11 @@ class JaxprEqnContext:
       yield
 
   def __repr__(self):
-    return (f"JaxprEqnContext(compute_type={self.compute_type},"
-            f"threefry_partitionable={self.threefry_partitionable})")
+    return (
+        f"JaxprEqnContext(compute_type={self.compute_type},"
+        f"threefry_partitionable={self.threefry_partitionable}),"
+        f"xla_metadata={self.xla_metadata}"
+    )
 
 
 class JaxprEqn:
@@ -333,8 +340,11 @@ class JaxprEqn:
 def new_jaxpr_eqn(invars, outvars, primitive, params, effects, source_info=None,
                   ctx=None):
   source_info = source_info or source_info_util.new_source_info()
-  ctx = ctx or JaxprEqnContext(compute_on.current_compute_type(),
-                               config.threefry_partitionable.value)
+  ctx = ctx or JaxprEqnContext(
+      compute_on.current_compute_type(),
+      config.threefry_partitionable.value,
+      xla_metadata_lib.current_xla_metadata(),
+  )
   if config.enable_checks.value:
     assert all(isinstance(x, (Var, Literal)) for x in  invars)
     assert all(isinstance(v,  Var)           for v in outvars)
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index 0e7e0146e..c4c77c72b 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -1923,6 +1923,10 @@ def lower_per_platform(ctx: LoweringRuleContext,
         lambda o: wrap_compute_type_in_place(ctx, o.owner),
         filter(_is_not_block_argument, flatten_ir_values(output)),
     )
+    map(
+        lambda o: wrap_xla_metadata_in_place(ctx, o.owner),
+        flatten_ir_values(output),
+    )
     return output
 
   assert len(platforms) > 1 and len(kept_rules) >= 2, (platforms, kept_rules)
@@ -1964,6 +1968,10 @@ def lower_per_platform(ctx: LoweringRuleContext,
           lambda o: wrap_compute_type_in_place(ctx, o.owner),
           filter(_is_not_block_argument, out_nodes),
       )
+      map(
+          lambda o: wrap_xla_metadata_in_place(ctx, o.owner),
+          out_nodes,
+      )
       if inner_ctx.tokens_out is not None:
         assert len(ordered_effects) == len(inner_ctx.tokens_out)
         out_nodes = [inner_ctx.tokens_out.get(eff)
@@ -2125,6 +2133,25 @@ def wrap_compute_type_in_place(ctx, op):
     op.operation.attributes["mhlo.frontend_attributes"] = ir.DictAttr.get(dict_attr)
 
 
+def wrap_xla_metadata_in_place(ctx, op):
+  ctx_attributes = {}
+  existing_attributes = {}
+  if ctx.jaxpr_eqn_ctx is not None and ctx.jaxpr_eqn_ctx.xla_metadata:
+    for k, v in ctx.jaxpr_eqn_ctx.xla_metadata.items():
+      ctx_attributes[k] = ir.StringAttr.get(str(v).lower())
+    if isinstance(op, ir.Operation):
+      # combine with existing mhlo.frontend_attributes
+      op_attributes_dict = {attr.name: attr.attr for attr in op.attributes}
+      for k, attributes in op_attributes_dict.items():
+        if k == "mhlo.frontend_attributes":
+          v_dict = {attr.name: attr.attr for attr in attributes}
+          for fa_key, fa_val in v_dict.items():
+            existing_attributes[fa_key] = fa_val
+      op.attributes["mhlo.frontend_attributes"] = ir.DictAttr.get(
+          ctx_attributes | existing_attributes
+      )
+
+
 def broadcast_in_dim(ctx: LoweringRuleContext, op, aval_out: core.AbstractValue, *,
                      broadcast_dimensions) -> ir.Value:
   # broadcast_dimension[i] is the axis of the result where the axis i of
diff --git a/jax/_src/interpreters/partial_eval.py b/jax/_src/interpreters/partial_eval.py
index 94f8918aa..5bb3e204c 100644
--- a/jax/_src/interpreters/partial_eval.py
+++ b/jax/_src/interpreters/partial_eval.py
@@ -35,6 +35,7 @@ from jax._src import linear_util as lu
 from jax._src import profiler
 from jax._src import source_info_util
 from jax._src import compute_on
+from jax._src import xla_metadata as xla_metadata_lib
 from jax._src.api_util import (flattened_fun_in_tree, flatten_fun_nokwargs,
                                fun_sourceinfo)
 from jax._src.core import (Trace, Tracer, Jaxpr, Literal, get_aval,
@@ -898,8 +899,11 @@ def new_eqn_recipe(in_tracers: Sequence[JaxprTracer],
     assert ("donated_invars" in params and
             len(params["donated_invars"]) == len(params["call_jaxpr"].invars))
   out_avals = [core.raise_to_shaped(t.aval) for t in out_tracers]
-  ctx = ctx or JaxprEqnContext(compute_on.current_compute_type(),
-                               config.threefry_partitionable.value)
+  ctx = ctx or JaxprEqnContext(
+      compute_on.current_compute_type(),
+      config.threefry_partitionable.value,
+      xla_metadata_lib.current_xla_metadata(),
+  )
   return JaxprEqnRecipe(object(), tuple(in_tracers), map(ref, out_tracers),
                         out_avals, primitive, params, effects, source_info,
                         ctx)
diff --git a/jax/_src/xla_metadata.py b/jax/_src/xla_metadata.py
new file mode 100644
index 000000000..94b482e2d
--- /dev/null
+++ b/jax/_src/xla_metadata.py
@@ -0,0 +1,55 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+import threading
+from contextlib import contextmanager
+
+from jax._src import config
+
+
+class _XlaMetadata(threading.local):
+  val: dict[Any, Any]
+
+  def __init__(self):
+    self.val = {}
+
+thread_local_metadata = _XlaMetadata()
+
+def current_xla_metadata():
+  return thread_local_metadata.val
+
+@contextmanager
+def set_xla_metadata(*args, **kwargs):
+  new_metadata = thread_local_metadata.val.copy()
+  if args:
+    new_metadata.update(args[0] if args[0] else {})
+  else:
+    new_metadata.update(**kwargs)
+  prev_metadata, thread_local_metadata.val = (
+      thread_local_metadata.val,
+      new_metadata,
+  )
+  config.update_thread_local_jit_state(
+      xla_metadata_context_manager=tuple(
+          (v, k) for k, v in sorted(new_metadata.items())))
+  try:
+    yield
+  finally:
+    thread_local_metadata.val = prev_metadata
+    config.update_thread_local_jit_state(
+        xla_metadata_context_manager=tuple(
+            (v, k) for k, v in sorted(prev_metadata.items())
+        )
+    )
diff --git a/jax/experimental/xla_metadata.py b/jax/experimental/xla_metadata.py
new file mode 100644
index 000000000..fb15e4743
--- /dev/null
+++ b/jax/experimental/xla_metadata.py
@@ -0,0 +1,17 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the ific language governing permissions and
+# limitations under the License.
+
+from jax._src.xla_metadata import (
+    set_xla_metadata as set_xla_metadata,
+)


Test patch:
diff --git a/tests/BUILD b/tests/BUILD
index b624b6bef..f8cf35c5d 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -1505,6 +1505,12 @@ jax_test(
     tags = ["multiaccelerator"],
 )
 
+jax_test(
+    name = "xla_metadata_test",
+    srcs = ["xla_metadata_test.py"],
+    deps = ["//jax:experimental"],
+)
+
 py_test(
     name = "pretty_printer_test",
     srcs = ["pretty_printer_test.py"],
diff --git a/tests/xla_metadata_test.py b/tests/xla_metadata_test.py
new file mode 100644
index 000000000..38bd7e055
--- /dev/null
+++ b/tests/xla_metadata_test.py
@@ -0,0 +1,290 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests whether the frontend attributes added by the context manager are
+
+correctly propagated to the jaxpr and mlir.
+"""
+
+from absl.testing import absltest
+import jax
+from jax._src import config
+from jax._src import dispatch
+from jax._src import test_util as jtu
+from jax._src.lax import lax
+from jax.experimental.xla_metadata import set_xla_metadata
+import jax.numpy as jnp
+
+config.parse_flags_with_absl()
+
+
+class XlaMetadataTest(jtu.JaxTestCase):
+
+  def test_f_jitted(self):
+    @jax.jit
+    def f(a, b):
+      with set_xla_metadata(a="b"):
+        return a + b
+
+    f_jaxpr = jax.make_jaxpr(f)(1, 2)
+    eqns = f_jaxpr.eqns
+    for eq in eqns[1:]:
+      self.assertDictEqual(eq.ctx.attributes, {"a": "b"})
+
+    f_lowered_text = f.lower(1.0, 2.0).as_text()
+    self.assertIn('mhlo.frontend_attributes = {a = "b"}', f_lowered_text)
+
+  def test_f_jitted_bool_attributes(self):
+    @jax.jit
+    def f(a, b):
+      with set_xla_metadata(a=True):
+        return a + b
+
+    f_lowered_text = f.lower(1.0, 2.0).as_text()
+    self.assertIn('mhlo.frontend_attributes = {a = "true"}', f_lowered_text)
+
+  def test_f_jitted_int_attributes(self):
+    @jax.jit
+    def f(a, b):
+      with set_xla_metadata(a=10):
+        return a + b
+
+    f_lowered_text = f.lower(1.0, 2.0).as_text()
+    self.assertIn('mhlo.frontend_attributes = {a = "10"}', f_lowered_text)
+
+  def test_f_nonjitted(self):
+    def f_add(a, b):
+      return dispatch.apply_primitive(lax.add_p, a, b)
+
+    arg1 = jnp.arange(2)
+    with set_xla_metadata(a="b"):
+      self.assertIn(
+          'mhlo.frontend_attributes = {a = "b"}',
+          jax.jit(f_add).lower(arg1, arg1).as_text(),
+      )
+
+  def test_f_attributes_overwrite(self):
+    @jax.jit
+    def g(a, b):
+      return a * b
+
+    with set_xla_metadata(a="b"):
+
+      @jax.jit
+      def f(a, b):
+        with set_xla_metadata(a="c"):
+          return a + b
+
+      f_lowered_text = f.lower(1.0, 2.0).as_text()
+      self.assertIn('mhlo.frontend_attributes = {a = "c"}', f_lowered_text)
+      self.assertIn(
+          'mhlo.frontend_attributes = {a = "b"}', g.lower(1.0, 2.0).as_text()
+      )
+    self.assertNotIn("mhlo.frontend_attributes", g.lower(1.0, 2.0).as_text())
+
+  def test_f_attributes_merge(self):
+    with set_xla_metadata(key1="val1"):
+
+      @jax.jit
+      def f(a, b):
+        with set_xla_metadata(key2="val2"):
+          return a + b
+
+      f_lowered_text = f.lower(1.0, 2.0).as_text()
+      self.assertIn(
+          'mhlo.frontend_attributes = {key1 = "val1", key2 = "val2"}',
+          f_lowered_text,
+      )
+
+  def test_attr_caching_jit(self):
+    @jax.jit
+    def f_add_jit(a, b):
+      return a + b
+
+    with set_xla_metadata(b="c"):
+      f_add_lowered1 = f_add_jit.lower(2.0, 3.0).as_text()
+    # Expect no attributes in the mlir.
+    f_add_lowered2 = f_add_jit.lower(1.0, 2.0).as_text()
+    with set_xla_metadata(c="d"):
+      f_add_lowered3 = f_add_jit.lower(4.0, 5.0).as_text()
+    self.assertIn('mhlo.frontend_attributes = {b = "c"}', f_add_lowered1)
+    self.assertNotIn("mhlo.frontend_attributes = {}", f_add_lowered2)
+    self.assertNotIn('mhlo.frontend_attributes = {b = "c"}', f_add_lowered2)
+    self.assertNotIn('mhlo.frontend_attributes = {c = "d"}', f_add_lowered2)
+    self.assertIn('mhlo.frontend_attributes = {c = "d"}', f_add_lowered3)
+
+  def test_attr_caching_nonjit(self):
+    def f_add(a, b):
+      return dispatch.apply_primitive(lax.add_p, a, b)
+
+    arg1 = jnp.arange(2)
+    arg2 = jnp.arange(2) + 1
+    arg3 = jnp.arange(2) + 2
+    with set_xla_metadata(b="c"):
+      self.assertIn(
+          'mhlo.frontend_attributes = {b = "c"}',
+          jax.jit(f_add).lower(arg1, arg1).as_text(),
+      )
+    # Expect no attributes in the jaxpr.
+    self.assertNotIn(
+        "mhlo.frontend_attributes",
+        jax.jit(f_add).lower(arg2, arg2).as_text(),
+    )
+
+    with set_xla_metadata(c="d"):
+      self.assertIn(
+          'mhlo.frontend_attributes = {c = "d"}',
+          jax.jit(f_add).lower(arg3, arg3).as_text(),
+      )
+
+  def test_axpy(self):
+    @jax.jit
+    def axpy(a, x, y):
+      with set_xla_metadata(a="b"):
+        return a * x + y
+
+    for line in axpy.lower(1.0, 2.0, 3.0).as_text().split("\n"):
+      if "stablehlo.multiply" in line:
+        self.assertIn('mhlo.frontend_attributes = {a = "b"}', line)
+      if "stablehlo.add" in line:
+        self.assertIn('mhlo.frontend_attributes = {a = "b"}', line)
+
+  def test_while(self):
+    @jax.jit
+    def f(a):
+      with set_xla_metadata(a="b"):
+        return jax.lax.while_loop(lambda x: x < 10, lambda x: x + 1, a)
+
+    self.assertIn(
+        'mhlo.frontend_attributes = {a = "b"}', f.lower(1.0).as_text()
+    )
+
+  def test_while_condition_body(self):
+    @jax.jit
+    def f_condition(x):
+      with set_xla_metadata(a="b"):
+        return x < 10
+
+    @jax.jit
+    def f_body(x):
+      with set_xla_metadata(a="c"):
+        return x + 1
+
+    @jax.jit
+    def while_fn(a):
+      return jax.lax.while_loop(f_condition, f_body, a)
+
+    for line in while_fn.lower(1.0).as_text().split("\n"):
+      if "stablehlo.compare" in line:
+        self.assertIn('mhlo.frontend_attributes = {a = "b"}', line)
+      if "stablehlo.add" in line:
+        self.assertIn('mhlo.frontend_attributes = {a = "c"}', line)
+
+  def test_nested_jit(self):
+    @jax.jit
+    def f(x, y):
+      with set_xla_metadata(a="b"):
+        z = x * y
+
+        @jax.jit
+        def g(z):
+          with set_xla_metadata(c="d"):
+            return z**2 + 1
+
+        return g(z)
+
+    self.assertIn(
+        'mhlo.frontend_attributes = {a = "b", c = "d"}',
+        f.lower(1.0, 2.0).as_text(),
+    )
+
+  def test_grad(self):
+    @jax.jit
+    def f(x, y):
+      with set_xla_metadata(a="b"):
+        return jax.grad(lambda x: x**3 + y**2 + jnp.sin(x))(x)
+
+    f_jaxpr = jax.make_jaxpr(f)(1.0, 2.0)
+    eqns = f_jaxpr.eqns
+    for eq in eqns[1:]:
+      self.assertDictEqual(eq.ctx.attributes, {"a": "b"})
+
+    self.assertIn(
+        'mhlo.frontend_attributes = {a = "b"}', f.lower(1.0, 2.).as_text()
+    )
+
+  def test_grad_outside_ctx(self):
+    @jax.jit
+    def f(x):
+      with set_xla_metadata(a="b"):
+        return x**3 + x**2 + jnp.sin(x)
+
+    grad_fn = jax.jit(jax.grad(f))
+    for line in grad_fn.lower(1.0).as_text().split("\n"):
+      if "stablehlo.cosine" in line:
+        self.assertIn('mhlo.frontend_attributes = {a = "b"}', line)
+      if "call @integer_pow" in line:
+        self.assertIn('mhlo.frontend_attributes = {a = "b"}', line)
+
+  def test_vmap(self):
+    dct = {"a": 0.0, "b": jnp.arange(5.0)}
+
+    @jax.jit
+    def f(dct, x):
+      with set_xla_metadata(a="b"):
+        return dct["a"] + dct["b"] + x
+
+    with set_xla_metadata(a="d"):
+      f_vmap = jax.vmap(f, in_axes=({"a": None, "b": 0}, None))
+      f_jaxpr = jax.make_jaxpr(f_vmap)(dct, 1.0)
+      eqns = f_jaxpr.eqns
+      for eq in eqns[1:]:
+        self.assertDictEqual(eq.ctx.attributes, {"a": "d"})
+    @jax.jit
+    def f2(x, y):
+      with set_xla_metadata(a="b"):
+        return (x + y, y * 2.0)
+
+    f_vmap_jaxpr = jax.make_jaxpr(jax.vmap(f2, in_axes=(0, None)))
+    self.assertIn(
+        'mhlo.frontend_attributes = {a = "b"}',
+        f_vmap_jaxpr.lower(jnp.arange(5.0), 1.0).as_text(),
+    )
+
+  def test_multiple_instructions(self):
+    @jax.jit
+    def f(x, a):
+      y = jnp.matmul(x, x)
+      with set_xla_metadata(a="b"):
+        return y + a
+
+    for line in f.lower(jnp.arange(5.0), 1.0).as_text().split("\n"):
+      # matmul doesn't have attributes
+      if "stablehlo.dot_general" in line:
+        self.assertNotIn('mhlo.frontend_attributes = {a = "b"}', line)
+      if "stablehlo.add" in line:
+        self.assertIn('mhlo.frontend_attributes = {a = "b"}', line)
+
+  def test_softmax(self):
+    @jax.jit
+    def f(x):
+      with set_xla_metadata(a="b"):
+        return jax.nn.softmax(x)
+    self.assertIn(
+        'mhlo.frontend_attributes = {a = "b"}', f.lower(jnp.arange(5.0)).as_text()
+    )
+
+
+if __name__ == "__main__":
+  absltest.main(testLoader=jtu.JaxTestLoader())