Repository: jax-ml/jax
Pull request: Change round/round_p API and implement jax2tf conversion & tests
Fixes google/jax#4952.

Problem statement:
[jax2tf] Incorrect translation for jax.lax.round
```
import jax
import jax.experimental.jax2tf as jax2tf
import jax.numpy as jnp
import tensorflow as tf
tf.compat.v1.enable_eager_execution()

print("Jax:", jax.lax.round(0.5))
print("Jax2Tf:", jax2tf.convert(jax.lax.round)(0.5))
```

```
Jax: 1.0
Jax2Tf: tf.Tensor(0.0, shape=(), dtype=float32)
```

The issue is coming from the fact that jax.lax.round uses away from zero rounding based on what the corresponding XLA op is doing but then it maps to tf,round what does round towards even.

Hints from issue discussion:
@bchetioui Is this now fixed?
@hawkinsp I think #4959 is ready, but it has not been merged!

Commit message: [jax2tf] Fix and test conversion of round.
Fixes google/jax#4952.

Patch to review:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index b8a298f71cff..8653b2fe7525 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -16,10 +16,12 @@
 # pytype: skip-file
 
 import builtins
+from enum import IntEnum
 import functools
 import itertools
 import operator
-from typing import (Any, Callable, List, NamedTuple, Optional, Sequence, Union, Tuple)
+from typing import (Any, Callable, List, NamedTuple, Optional, Sequence,\
+                    Union, Tuple)
 import warnings
 
 import numpy as np
@@ -150,12 +152,28 @@ def ceil(x: Array) -> Array:
   r"""Elementwise ceiling: :math:`\left\lceil x \right\rceil`."""
   return ceil_p.bind(x)
 
-def round(x: Array) -> Array:
+class RoundingMethod(IntEnum):
+  AWAY_FROM_ZERO = 0
+  TO_NEAREST_EVEN = 1
+
+def round(x: Array,
+          rounding_method: RoundingMethod = RoundingMethod.AWAY_FROM_ZERO
+          ) -> Array:
   r"""Elementwise round.
 
-  Rounds values to the nearest integer. Halfway values (e.g., `0.5`) are rounded
-  away from zero."""
-  return round_p.bind(x)
+  Rounds values to the nearest integer.
+
+  Args:
+    x: an array or scalar value to round.
+    rounding_method: the method to use when rounding halfway values
+      (e.g., `0.5`). See ``lax.RoundingMethod`` for the list of possible
+      values.
+
+  Returns:
+    An array containing the elementwise rounding of x.
+  """
+  rounding_method = RoundingMethod(rounding_method)
+  return round_p.bind(x, rounding_method=rounding_method)
 
 def is_finite(x: Array) -> Array:
   r"""Elementwise :math:`\mathrm{isfinite}`.
@@ -2167,7 +2185,28 @@ def _sign_translation_rule(c, x):
 ceil_p = standard_unop(_float, 'ceil')
 ad.defjvp_zero(ceil_p)
 
+def _round_to_nearest_even(x):
+  half = _const(x, 0.5)
+  one = _const(x, 1)
+  round_val = floor(x)
+  fraction = x - round_val
+  nearest_even_int = sub(
+    round_val, mul(_const(x, 2), floor(mul(half, x))))
+  is_odd = eq(nearest_even_int, one)
+  return select(
+    bitwise_or(gt(fraction, half),
+               bitwise_and(eq(fraction, half), is_odd)),
+    add(round_val, one), round_val)
+
+def _round_translation_rule(c, x, *, rounding_method):
+  if rounding_method is RoundingMethod.AWAY_FROM_ZERO:
+    return xops.Round(x)
+  else: # rounding_method is RoundingMethod.TO_NEAREST_EVEN
+    rounding_fun = xla.lower_fun(_round_to_nearest_even, multiple_results=False)
+    return rounding_fun(c, x)
+
 round_p = standard_unop(_float, 'round')
+xla.translations[round_p] = _round_translation_rule
 ad.defjvp_zero(round_p)
 
 is_finite_p = unop(_fixed_dtype(np.bool_), _float, 'is_finite')
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index f7b5cfae7d66..9eb7465b6341 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -495,7 +495,7 @@ def rint(x):
     return lax.convert_element_type(x, float_)
   if issubdtype(dtype, complexfloating):
     return lax.complex(rint(lax.real(x)), rint(lax.imag(x)))
-  return _round_to_nearest_even(x)
+  return lax.round(x, lax.RoundingMethod.TO_NEAREST_EVEN)
 
 
 @_wraps(np.sign)
@@ -1695,20 +1695,6 @@ def clip(a, a_min=None, a_max=None, out=None):
     a = minimum(a_max, a)
   return a
 
-
-def _round_to_nearest_even(x):
-  half = lax._const(x, 0.5)
-  one = lax._const(x, 1)
-  round_val = lax.floor(x)
-  fraction = x - round_val
-  nearest_even_int = lax.sub(
-    round_val, lax.mul(lax._const(x, 2), lax.floor(lax.mul(half, x))))
-  is_odd = lax.eq(nearest_even_int, one)
-  return lax.select(
-    lax.bitwise_or(lax.gt(fraction, half),
-                   lax.bitwise_and(lax.eq(fraction, half), is_odd)),
-    lax.add(round_val, one), round_val)
-
 @_wraps(np.round, update_doc=False)
 def round(a, decimals=0, out=None):
   _check_arraylike("round", a)
@@ -1723,7 +1709,7 @@ def round(a, decimals=0, out=None):
 
   def _round_float(x):
     if decimals == 0:
-      return _round_to_nearest_even(x)
+      return lax.round(x, lax.RoundingMethod.TO_NEAREST_EVEN)
 
     # TODO(phawkins): the strategy of rescaling the value isn't necessarily a
     # good one since we may be left with an incorrectly rounded value at the
@@ -1731,7 +1717,8 @@ def _round_float(x):
     # float32,
     x = lax.convert_element_type(x, np.float32) if dtype == np.float16 else x
     factor = _constant_like(x, 10 ** decimals)
-    out = lax.div(_round_to_nearest_even(lax.mul(x, factor)), factor)
+    out = lax.div(lax.round(lax.mul(x, factor),
+                            lax.RoundingMethod.TO_NEAREST_EVEN), factor)
     return lax.convert_element_type(out, dtype) if dtype == np.float16 else out
 
   if issubdtype(dtype, complexfloating):
diff --git a/jax/experimental/jax2tf/jax2tf.py b/jax/experimental/jax2tf/jax2tf.py
index 6515c41ab904..3cb3ee3067b0 100644
--- a/jax/experimental/jax2tf/jax2tf.py
+++ b/jax/experimental/jax2tf/jax2tf.py
@@ -854,7 +854,20 @@ def _unexpected_primitive(p: core.Primitive, *args, **kwargs):
 tf_impl[lax.sign_p] = tf.math.sign
 tf_impl[lax.floor_p] = tf.math.floor
 tf_impl[lax.ceil_p] = tf.math.ceil
-tf_impl[lax.round_p] = tf.math.round
+
+def _round(operand, *, rounding_method):
+  if rounding_method is lax.RoundingMethod.AWAY_FROM_ZERO:
+    sign = tf.math.sign(operand)
+    operand *= sign
+    floor = tf.math.floor(operand)
+    operand -= floor
+    cond = tf.math.equal(operand, tf.constant(np.array(0.5), operand.dtype))
+    return sign * (tf.where(cond, tf.constant(np.array(1), operand.dtype),
+                            tf.math.round(operand)) + floor)
+  else:
+    return tf.math.round(operand)
+
+tf_impl[lax.round_p] = _round
 tf_impl[lax.nextafter_p] = tf.math.nextafter
 
 def _population_count(x):
diff --git a/jax/experimental/jax2tf/tests/primitive_harness.py b/jax/experimental/jax2tf/tests/primitive_harness.py
index d1a0752900b1..da27e5e96222 100644
--- a/jax/experimental/jax2tf/tests/primitive_harness.py
+++ b/jax/experimental/jax2tf/tests/primitive_harness.py
@@ -179,6 +179,37 @@ def parameterized(harness_group: Iterable[Harness],
   ]
 )
 
+def _make_round_harness(name, *, shape=(100, 100), dtype=np.float32,
+                        rounding_method=lax.RoundingMethod.AWAY_FROM_ZERO,
+                        operand=None):
+  operand = operand if operand is not None else RandArg(shape, dtype)
+  return Harness(f"{name}_shape={jtu.format_shape_dtype_string(operand.shape, operand.dtype)}_roundingmethod={rounding_method}",
+                 lax.round,
+                 [operand, StaticArg(rounding_method)],
+                 operand=operand)
+
+lax_round = tuple( # Validate dtypes
+  _make_round_harness("dtypes", dtype=dtype)
+  for dtype in jtu.dtypes.all_floating
+) + tuple( # Validate rounding method
+  _make_round_harness("dtypes", operand=operand,
+                      rounding_method=rounding_method)
+  for operand in [
+    np.array([[0.5, 1.5, 2.5], [-0.5, -1.5, -2.5]], dtype=np.float32)
+  ]
+  for rounding_method in [
+    lax.RoundingMethod.AWAY_FROM_ZERO,
+    lax.RoundingMethod.TO_NEAREST_EVEN
+  ]
+) + tuple( # Validate edge cases
+  _make_round_harness(f"edge_case_{name}", operand=operand)
+  for name, operand in [
+    # Checks that https://github.com/google/jax/issues/4952 is resolved
+    ("round_away_from_0", np.array([[0.5, 1.5, 2.5],
+                                    [-0.5, -1.5, -2.5]], dtype=np.float32)),
+  ]
+)
+
 _LAX_COMPARATORS = (
   lax.eq, lax.ge, lax.gt, lax.le, lax.lt, lax.ne)
 
diff --git a/jax/experimental/jax2tf/tests/primitives_test.py b/jax/experimental/jax2tf/tests/primitives_test.py
index ceaaa9e183c0..a8656f2d70d7 100644
--- a/jax/experimental/jax2tf/tests/primitives_test.py
+++ b/jax/experimental/jax2tf/tests/primitives_test.py
@@ -570,6 +570,10 @@ def custom_assert(result_jax, result_tf):  # noqa: F811
     self.ConvertAndCompare(harness.dyn_fun, arg, custom_assert=custom_assert,
                            atol=atol)
 
+  @primitive_harness.parameterized(primitive_harness.lax_round)
+  def test_round(self, harness: primitive_harness.Harness):
+    self.ConvertAndCompare(harness.dyn_fun, *harness.dyn_args_maker(self.rng()))
+
   @primitive_harness.parameterized(primitive_harness.lax_comparators)
   def test_comparators(self, harness: primitive_harness.Harness):
     self.ConvertAndCompare(harness.dyn_fun, *harness.dyn_args_maker(self.rng()))
diff --git a/jax/lax/__init__.py b/jax/lax/__init__.py
index 013766d61eb6..d0d0f08ec313 100644
--- a/jax/lax/__init__.py
+++ b/jax/lax/__init__.py
@@ -19,6 +19,7 @@
   DotDimensionNumbers,
   GatherDimensionNumbers,
   Precision,
+  RoundingMethod,
   ScatterDimensionNumbers,
   abs,
   abs_p,
diff --git a/tests/lax_test.py b/tests/lax_test.py
index f157785bf428..2d4e02f89dd8 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1177,6 +1177,13 @@ def testReshapeAgainstNumpy(self, arg_shape, out_shape, dtype, rng_factory):
     numpy_op = lambda x: lax_reference.reshape(x, out_shape)
     self._CheckAgainstNumpy(numpy_op, op, args_maker)
 
+  def testRoundRoundingMethods(self):
+    x = np.array([-2.5, -1.5, -0.5, 0.5, 1.5, 2.5], dtype=np.float32)
+    self.assertAllClose(lax.round(x, lax.RoundingMethod.AWAY_FROM_ZERO),
+                        np.array([-3, -2, -1, 1, 2, 3], dtype=np.float32))
+    self.assertAllClose(lax.round(x, lax.RoundingMethod.TO_NEAREST_EVEN),
+                        np.array([-2, -2, 0, 0, 2, 2], dtype=np.float32))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_inshape={}_pads={}"
        .format(jtu.format_shape_dtype_string(shape, dtype), pads),


Code review comments:
diff_hunk: @@ -196,3 +196,14 @@ Linear algebra operators (jax.lax.linalg)
     qr
     svd
     triangular_solve
+
+Argument wrappers
path: docs/jax.lax.rst
text: @user1:
I might call this "Types" or "Classes" or something like that. I don't find "wrappers" very descriptive.