Fix comparator name.

hawkinsp · Jul 1, 2020 · 89a860b · 89a860b
1 parent 8276d46
commit 89a860b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
@@ -4322,7 +4322,7 @@ def _argminmax_translation_rule(value_comparator, identity,
   dtype = shape.numpy_dtype()
   index_dtype = _int_dtype_for_dim_index(shape.dimensions()[axis])
 
-  subc = xb.make_computation_builder("argmax_comparator")
+  subc = xb.make_computation_builder("argminmax_comparator")
   value_shape = xc.Shape.array_shape(shape.xla_element_type(), ())
   index_shape = xc.Shape.array_shape(index_dtype, ())
   x_value = xb.parameter(subc, 0, value_shape)