Repository: google/jax

Issue:
Add lax.reduce_precision primitive corresponding to XLA's ReducePrecision op
I want to round a float32 tensor x to bfloat16 values and get it back as float32.
The natural thing is: x.astype(jnp.bfloat32).astype(jnp.float32)
But it looks like XLA optimizes this away to just x again?
Could we have an op to enforce bfloat16 rounding? (e.g., to simulate bfloat16 training in parts of the net)


Discussion:
It looks like your strategy works, at least in simple cases:
```python
import jax.numpy as jnp
from jax import jit

@jit
def round(x):
  return x.astype(jnp.bfloat16).astype(jnp.float32)

x = jnp.float32(jnp.pi)
y = round(x)
print(x, y, x == y)
# 3.1415927 3.140625 False
```
Do you have an example of a place where this approach fails?
XLA is allowed to make optimizations that increase precision. I know the TPU backend does so in some cases, if nothing else. The only way to guarantee precision is not increased by XLA would be to add a JAX primitive that wraps the `ReducePrecision` operator, that exists explicitly for this purpose.
I'll take a look at exposing this in `lax.py`

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -5449,6 +5449,34 @@ def reduce_window_shape_tuple(operand_shape, window_dimensions, window_strides,
   _reduce_window_batch_rule, _reduce_window_min)
 
 
+def _reduce_precision_shape_rule(operand, *, exponent_bits, mantissa_bits):
+  exponent_bits = operator.index(exponent_bits)
+  mantissa_bits = operator.index(mantissa_bits)
+  if exponent_bits < 1:
+    raise ValueError(f"reduce_precision: exponent_bits must be positive; got {exponent_bits}")
+  if mantissa_bits < 0:
+    raise ValueError(f"reduce_precision: mantissa_bits must be non-negative; got {mantissa_bits}")
+  return operand.shape
+
+
+reduce_precision_p = standard_primitive(
+    _reduce_precision_shape_rule,
+    partial(unop_dtype_rule, _identity, _float, 'reduce_precision'),
+    name='reduce_precision')
+
+
+def reduce_precision(operand, exponent_bits, mantissa_bits):
+  """Wraps XLA's `ReducePrecision
+  <https://www.tensorflow.org/xla/operation_semantics#reduceprecision>`_
+  operator.
+  """
+  exponent_bits = core.concrete_or_error(
+    operator.index, exponent_bits, "exponent_bits argument of lax.reduce_precision")
+  mantissa_bits = core.concrete_or_error(
+    operator.index, mantissa_bits, "mantissa_bits argument of lax.reduce_precision")
+  return reduce_precision_p.bind(operand, exponent_bits=exponent_bits, mantissa_bits=mantissa_bits)
+
+
 def _select_and_scatter_shape_rule(
     operand, source, init_value, *, select_jaxpr, select_consts, scatter_jaxpr,
     scatter_consts, window_dimensions, window_strides, padding):
diff --git a/jax/experimental/jax2tf/jax2tf.py b/jax/experimental/jax2tf/jax2tf.py
--- a/jax/experimental/jax2tf/jax2tf.py
+++ b/jax/experimental/jax2tf/jax2tf.py
@@ -794,6 +794,7 @@ def _unexpected_primitive(p: core.Primitive, *args, **kwargs):
 
   "igamma_grad_a",
   "random_gamma_grad",
+  "reduce_precision",
 
   # Not high priority?
   "after_all", "all_to_all", "create_token",
diff --git a/jax/lax/__init__.py b/jax/lax/__init__.py
--- a/jax/lax/__init__.py
+++ b/jax/lax/__init__.py
@@ -209,6 +209,8 @@
   reduce_min_p,
   reduce_or_p,
   reduce_p,
+  reduce_precision,
+  reduce_precision_p,
   reduce_prod_p,
   reduce_sum_p,
   reduce_window,


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1693,6 +1693,25 @@ def np_fun(x):
     self._CompileAndCheck(fun, args_maker)
     self._CheckAgainstNumpy(np_fun, fun, args_maker)
 
+
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_shape={}_out_dtype={}".format(
+          jtu.format_shape_dtype_string(shape, dtype),
+          jtu.format_shape_dtype_string(shape, out_dtype)),
+       "shape": shape, "dtype": dtype, "out_dtype": out_dtype}
+      for shape in [(), (3,), (3, 4)]
+      for dtype in float_dtypes
+      for out_dtype in float_dtypes))
+  def testReducePrecision(self, shape, dtype, out_dtype):
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: [rng(shape, dtype)]
+    info = dtypes.finfo(out_dtype)
+    fun = lambda x: lax.reduce_precision(x, info.nexp, info.nmant)
+    np_fun = lambda x: np.asarray(x).astype(out_dtype).astype(dtype)
+    self._CheckAgainstNumpy(np_fun, fun, args_maker)
+    self._CompileAndCheck(fun, args_maker)
+
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_axis={}_isstable={}".format(
           jtu.format_shape_dtype_string(shape, dtype), axis, is_stable),