Repository: google/jax
Pull request #19537 (2024-01-26 23:42:13)

Avoid using lambda as the reducer fn
Avoid using lambda as the reducer fn

Lambdas are represented by their ids in the metadata of lowered HLO (see example below) and they change every time. This makes the compilation cache less effective as it causes the computation's fingerprint to change every time.

```
get-tuple-element.41724 = bf16[8]{0} get-tuple-element(reduce.41723), index=0, metadata={op_name="pjit(_wrapped_fn)/jit(main)/.../reduce[computation=<function _compute_argminmax.<locals>.reducer_fn at 0x7fa6ecfb2200> dimensions=(1,)]" source_file="..." source_line=...}
```



Patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 5480f61cf..4ad39a398 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -2046,8 +2046,6 @@ def _pow_lower(ctx, x, y):
   return _nary_lower_hlo(hlo.power, ctx_, x_, y_)
 mlir.register_lowering(pow_p, _pow_lower)
 
-
-
 def _integer_pow_dtype_rule(x, *, y):
   dtype = unop_dtype_rule(_identity, _int | _float | _complex, 'integer_pow', x)
   if y < 0 and dtypes.issubdtype(dtype, np.integer):
@@ -3928,17 +3926,22 @@ def _argminmax_dtype_rule(operand, *, axes, index_dtype):
                     .format(dtype_to_string(index_dtype)))
   return index_dtype
 
-def _compute_argminmax(value_comparator, get_identity,
-                       operand, *, index_dtype, axes):
-  # value_comparator is either lax.lt (for argmin) or lax.gt
-  # get_identity(operand.dtype) is inf for argmin or -inf for argmax
-  axis, = axes
-  indices = broadcasted_iota(index_dtype, np.shape(operand), axis)
-  def reducer_fn(op_val_index, acc_val_index):
+class _ArgMinMaxReducer:
+
+  def __init__(self, value_comparator):
+    self._value_comparator = value_comparator
+
+  def __repr__(self):
+    # Override the repr so that the metadata attached to the lowered op does not
+    # contain unstable function ids. This plays more nicely with computation
+    # fingerprint calculation in the compilation cache.
+    return f'_ArgMinMaxReducer({self._value_comparator.__name__})'
+
+  def __call__(self, op_val_index, acc_val_index):
     op_val, op_index = op_val_index
     acc_val, acc_index = acc_val_index
     # Pick op_val if Lt (for argmin) or if NaN
-    pick_op_val = bitwise_or(value_comparator(op_val, acc_val),
+    pick_op_val = bitwise_or(self._value_comparator(op_val, acc_val),
                              ne(op_val, op_val))
     # If x and y are not NaN and x = y, then pick the first
     pick_op_index = bitwise_or(pick_op_val,
@@ -3946,9 +3949,16 @@ def _compute_argminmax(value_comparator, get_identity,
                                            lt(op_index, acc_index)))
     return (select(pick_op_val, op_val, acc_val),
             select(pick_op_index, op_index, acc_index))
+
+def _compute_argminmax(value_comparator, get_identity,
+                       operand, *, index_dtype, axes):
+  # value_comparator is either lax.lt (for argmin) or lax.gt
+  # get_identity(operand.dtype) is inf for argmin or -inf for argmax
+  axis, = axes
+  indices = broadcasted_iota(index_dtype, np.shape(operand), axis)
   res = reduce([operand, indices],
                [get_identity(operand.dtype), np.array(0, index_dtype)],
-               reducer_fn,
+               _ArgMinMaxReducer(value_comparator),
                axes)
   return res[1]
 


Test patch:
diff --git a/tests/metadata_test.py b/tests/metadata_test.py
index 86a85046d..3511595d6 100644
--- a/tests/metadata_test.py
+++ b/tests/metadata_test.py
@@ -86,6 +86,12 @@ class MetadataTest(jtu.JaxTestCase):
     self.assertRegex(hlo, r'loc\(".*cond/branch_0_fun/cos"')
     self.assertRegex(hlo, r'loc\(".*cond/branch_1_fun/sin"')
 
+  def test_argmax(self):
+    def f(x):
+      return jnp.argmax(x)
+    hlo = module_to_string(jax.jit(f).lower(jnp.arange(8.0)).compiler_ir())
+    self.assertNotRegex(hlo, r'<.* at 0x[0-9a-fA-F]+>')
+
   def test_source_file_prefix_removal(self):
 
     def make_hlo():