add dot precision flag, enum state utility

Author: mattjj

jax/__init__.py

@@ -31,7 +31,8 @@
 
 # flake8: noqa: F401
 from .config import (config, enable_checks, check_tracer_leaks, checking_leaks,
-                     debug_nans, debug_infs, log_compiles)
+                     debug_nans, debug_infs, log_compiles,
+                     default_dot_precision, numpy_rank_promotion)
 from .api import (
   ad,  # TODO(phawkins): update users to avoid this.
   argnums_partial,  # TODO(phawkins): update Haiku to not use this.

jax/_src/lax/lax.py

@@ -496,7 +496,13 @@ def concatenate(operands: Sequence[Array], dimension: int) -> Array:
 Precision = xla_client.PrecisionConfig.Precision
 Precision.__str__ = lambda precision: precision.name
 PrecisionType = Any
-PrecisionLike = Union[None, PrecisionType, Tuple[PrecisionType, PrecisionType]]
+PrecisionLike = Union[None, str, PrecisionType,
+                      Tuple[PrecisionType, PrecisionType]]
+_precision_strings = {
+    'bfloat16':      Precision.DEFAULT,
+    'tensorfloat32': Precision.HIGH,
+    'float32':       Precision.HIGHEST,
+}
 
 
 class ConvDimensionNumbers(NamedTuple):
@@ -551,9 +557,10 @@ def conv_general_dilated(
     feature_group_count: integer, default 1. See XLA HLO docs.
     batch_group_count: integer, default 1. See XLA HLO docs.
     precision: Optional. Either ``None``, which means the default precision for
-      the backend, a ``lax.Precision`` enum value (``Precision.DEFAULT``,
-      ``Precision.HIGH`` or ``Precision.HIGHEST``) or a tuple of two
-      ``lax.Precision`` enums indicating precision of ``lhs``` and ``rhs``.
+      the backend, a string ('bfloat16', 'tensorfloat32', or 'float32'), a
+      ``lax.Precision`` enum value (``Precision.DEFAULT``, ``Precision.HIGH`` or
+      ``Precision.HIGHEST``) or a tuple of two ``lax.Precision`` enums
+      indicating precision of ``lhs``` and ``rhs``.
 
   Returns:
     An array containing the convolution result.
@@ -6378,16 +6385,19 @@ def remaining(original, *removed_lists):
 
 def _canonicalize_precision(precision):
   if precision is None:
-    return None
-  if isinstance(precision, Precision) or (
-      isinstance(precision, tuple)
-      and len(precision) == 2
-      and all(isinstance(p, Precision) for p in precision)
-  ):
+    return _precision_strings.get(config.jax_default_dot_precision)
+  elif isinstance(precision, str) and precision in _precision_strings:
+    return _precision_strings.get(precision)
+  elif isinstance(precision, Precision):
+    return precision
+  elif (isinstance(precision, (list, tuple)) and len(precision) == 2 and
+        all(isinstance(p, Precision) for p in precision)):
     return precision
   else:
-    raise ValueError("Precision argument must be None, a lax.Precision value "
-                     f"or a tuple of two lax.Precision values; got {precision}")
+    raise ValueError(
+        f"Precision argument must be None, a string in {_precision_strings}, "
+        "a lax.Precision value or a tuple of two lax.Precision values; "
+        f"got {precision}.")
 
 
 def conv_dimension_numbers(lhs_shape, rhs_shape, dimension_numbers

jax/_src/numpy/lax_numpy.py

@@ -54,14 +54,6 @@
                            canonicalize_axis as _canonicalize_axis, maybe_named_axis)
 from jax.tree_util import tree_leaves, tree_flatten, tree_map
 
-FLAGS = flags.FLAGS
-flags.DEFINE_enum(
-    'jax_numpy_rank_promotion', os.getenv('JAX_NUMPY_RANK_PROMOTION', 'allow'),
-    enum_values=['allow', 'warn', 'raise'],
-    help=
-    'Control NumPy-style automatic rank promotion broadcasting '
-    '("allow", "warn", or "raise").')
-
 newaxis = None
 
 # Common docstring additions:
@@ -246,20 +238,20 @@ def _promote_shapes(fun_name, *args):
     if not nonscalar_ranks or len(set(nonscalar_ranks)) == 1:
       return args
     else:
-      if FLAGS.jax_numpy_rank_promotion != "allow":
+      if config.jax_numpy_rank_promotion != "allow":
         _rank_promotion_warning_or_error(fun_name, shapes)
       result_rank = len(lax.broadcast_shapes(*shapes))
       return [broadcast_to(arg, (1,) * (result_rank - len(shp)) + shp)
               for arg, shp in zip(args, shapes)]
 
 def _rank_promotion_warning_or_error(fun_name, shapes):
-  if FLAGS.jax_numpy_rank_promotion == "warn":
+  if config.jax_numpy_rank_promotion == "warn":
     msg = ("Following NumPy automatic rank promotion for {} on shapes {}. "
            "Set the jax_numpy_rank_promotion config option to 'allow' to "
            "disable this warning; for more information, see "
            "https://jax.readthedocs.io/en/latest/rank_promotion_warning.html.")
     warnings.warn(msg.format(fun_name, ' '.join(map(str, shapes))))
-  elif FLAGS.jax_numpy_rank_promotion == "raise":
+  elif config.jax_numpy_rank_promotion == "raise":
     msg = ("Operands could not be broadcast together for {} on shapes {} "
            "and with the config option jax_numpy_rank_promotion='raise'. "
            "For more information, see "

jax/api.py

@@ -356,7 +356,9 @@ def f_jitted(*args, **kwargs):
     # TODO(jblespiau): We can remove `config.x64_enabled` when jaxlib 0.1.62 is
     # the minimal version. NOTE(mattjj): minversion 0.1.62 didn't work...
     context = (getattr(core.thread_local_state.trace_state.trace_stack,
-                       "dynamic", None), config.x64_enabled)
+                       "dynamic", None),
+               config.x64_enabled,
+               config.jax_default_dot_precision)
     # TODO(jblespiau): Move this to C++.
     if (config.jax_debug_nans or config.jax_debug_infs) and not _jit_is_disabled():
       device_arrays = cpp_jitted_f(context, *args, **kwargs)
@@ -2442,6 +2444,21 @@ def named_f(*args, **kwargs):
 
   return named_f
 
+
+def invertible(fun: Callable) -> Callable:
+  """Asserts that the decorated function is invertible.
+
+  Applying reverse-mode AD to a decorated function will use a more memory efficient
+  procedure than usual, which will reconstruct the necessary intermediate values
+  by inverting the function. Note that this might degrade the numerical accuracy of
+  obtained gradients if the inverse is unstable.
+
+  Args:
+    fun: The function assumed to be invertible.
+  """
+  return iad.invertible(fun)
+
+
 # TODO(mattjj): delete everything below here (deprecated custom_transforms)
 
 class CustomTransformsFunction(object):
@@ -2583,16 +2600,3 @@ def vjpfun(ct):
                    for x, vjp in zip(primals, vjprules))
     return ans, vjpfun
   defvjp_all(fun, custom_vjp)
-
-def invertible(fun: Callable) -> Callable:
-  """Asserts that the decorated function is invertible.
-
-  Applying reverse-mode AD to a decorated function will use a more memory efficient
-  procedure than usual, which will reconstruct the necessary intermediate values
-  by inverting the function. Note that this might degrade the numerical accuracy of
-  obtained gradients if the inverse is unstable.
-
-  Args:
-    fun: The function assumed to be invertible.
-  """
-  return iad.invertible(fun)

jax/config.py

@@ -1,6 +1,6 @@
-# Copyright 2018 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
+ # Copyright 2018 Google LLC
+ #
+ # 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
 #
@@ -17,6 +17,7 @@
 import os
 import sys
 import threading
+from typing import List, Optional
 
 from jax import lib
 
@@ -51,7 +52,7 @@ class Config:
   def __init__(self):
     self.values = {}
     self.meta = {}
-    self.FLAGS = NameSpace(self.read)
+    self.FLAGS = NameSpace(self.read, self.update)
     self.use_absl = False
     self._contextmanager_flags = set()
 
@@ -219,6 +220,54 @@ def get_state(self):
     @contextlib.contextmanager
     def set_state(new_val: bool):
       prev_val = getattr(_thread_local_state, name, unset)
+      setattr(_thread_local_state, name, bool(new_val))
+      try:
+        yield
+      finally:
+        if prev_val is unset:
+          delattr(_thread_local_state, name)
+        else:
+          setattr(_thread_local_state, name, prev_val)
+    set_state.__name__ = name[4:] if name.startswith('jax_') else name
+    set_state.__doc__ = f"Context manager for `{name}` config option.\n\n{help}"
+    return set_state
+
+  def define_enum_state(self, name: str, enum_values: List[str],
+                        default: Optional[str], help: str):
+    """Set up thread-local state and return a contextmanager for managing it.
+
+    Args:
+      name: string, converted to lowercase to define the name of the config
+        option (and absl flag). It is converted to uppercase to define the
+        corresponding shell environment variable.
+      enum_values: list of strings representing the possible values for the
+        option.
+      default: optional string, default value.
+      help: string, used to populate the flag help information as well as the
+        docstring of the returned context manager.
+
+    Returns:
+      A contextmanager to control the thread-local state value.
+
+    See docstring for ``define_bool_state``.
+    """
+    name = name.lower()
+    self.DEFINE_enum(name, os.getenv(name.upper(), default),
+                     enum_values=enum_values, help=help)
+    self._contextmanager_flags.add(name)
+
+    def get_state(self):
+      val = getattr(_thread_local_state, name, unset)
+      return val if val is not unset else self._read(name)
+    setattr(Config, name, property(get_state))
+
+    @contextlib.contextmanager
+    def set_state(new_val: Optional[str]):
+      if (new_val is not None and
+          (type(new_val) is not str or new_val not in enum_values)):
+        raise ValueError(f"new enum value must be None or in {enum_values}, "
+                         f"got {new_val} of type {type(new_val)}.")
+      prev_val = getattr(_thread_local_state, name, unset)
       setattr(_thread_local_state, name, new_val)
       try:
         yield
@@ -231,18 +280,25 @@ def set_state(new_val: bool):
     set_state.__doc__ = f"Context manager for `{name}` config option.\n\n{help}"
     return set_state
 
+
 _thread_local_state = threading.local()
 
 class Unset: pass
 unset = Unset()
 
-class NameSpace(object):
-  def __init__(self, getter):
-    self._getter = getter
+class NameSpace:
+  def __init__(self, getter, setter):
+    # must use super because we override this class's __setattr__, see
+    # https://docs.python.org/3/reference/datamodel.html#object.__setattr__
+    super().__setattr__('_getter', getter)
+    super().__setattr__('_setter', setter)
 
   def __getattr__(self, name):
     return self._getter(name)
 
+  def __setattr__(self, name, val):
+    self._setter(name, val)
+
 
 config = Config()
 flags = config
@@ -316,3 +372,32 @@ def __getattr__(self, name):
           'computation. Logging is performed with `absl.logging`. When this '
           'option is set, the log level is WARNING; otherwise the level is '
           'DEBUG.'))
+
+
+numpy_rank_promotion = config.define_enum_state(
+    name='jax_numpy_rank_promotion',
+    enum_values=['allow', 'warn', 'raise'],
+    default='allow',
+    help=('Control NumPy-style automatic rank promotion broadcasting '
+          '("allow", "warn", or "raise").'))
+
+default_dot_precision = config.define_enum_state(
+    name='jax_default_dot_precision',
+    enum_values=['bfloat16', 'tensorfloat32', 'float32'],
+    default=None,
+    help=('Control the default matmul and conv precision for 32bit inputs.\n\n'
+
+          'Some platforms, like TPU, offer configurable precision levels for '
+          'matrix multiplication and convolution computations, trading off '
+          'accuracy for speed. The precision can be controlled for each '
+          'operation; for example, see the :func:`jax.lax.conv_general_dilated` '
+          'and :func:`jax.lax.dot` docstrings. But it can be useful to control '
+          'the default behavior obtained when an operation is not given a '
+          'specific precision.\n\n'
+
+          'This option can be used to control the default precision '
+          'level for computations involved in matrix multiplication and '
+          'convolution on 32bit inputs. The levels roughly describe the '
+          "precision at which scalar products are computed. The 'bfloat16' "
+          "option is the fastest and least precise; 'float32' is similar to "
+          "full 'float32 precision; 'tensorfloat32' is intermediate.\n\n"))

tests/api_test.py

@@ -2393,6 +2393,58 @@ def f(_):
     expected = jnp.arange(1) + 1
     self.assertAllClose(ans, expected)
 
+  def test_dot_precision_context_manager(self):
+    x = jnp.zeros((2, 2))
+
+    with jax.default_dot_precision(None):
+      jnp.dot(x, x)  # doesn't crash
+      jaxpr = jax.make_jaxpr(jnp.dot)(x, x)
+    self.assertIn('precision=None', str(jaxpr))
+
+    with jax.default_dot_precision("bfloat16"):
+      jnp.dot(x, x)  # doesn't crash
+      jaxpr = jax.make_jaxpr(jnp.dot)(x, x)
+    self.assertIn('precision=DEFAULT', str(jaxpr))
+
+    with jax.default_dot_precision("tensorfloat32"):
+      jnp.dot(x, x)  # doesn't crash
+      jaxpr = jax.make_jaxpr(jnp.dot)(x, x)
+    self.assertIn('precision=HIGH', str(jaxpr))
+
+    with jax.default_dot_precision("float32"):
+      jnp.dot(x, x)  # doesn't crash
+      jaxpr = jax.make_jaxpr(jnp.dot)(x, x)
+    self.assertIn('precision=HIGHEST', str(jaxpr))
+
+    dot = partial(jnp.dot, precision=lax.Precision.HIGHEST)
+    with jax.default_dot_precision("tensorfloat32"):
+      dot(x, x)  # doesn't crash
+      jaxpr = jax.make_jaxpr(dot)(x, x)
+    self.assertIn('precision=HIGHEST', str(jaxpr))
+
+  def test_dot_precision_flag(self):
+    x = jnp.zeros((2, 2))
+
+    prev_val = config._read("jax_default_dot_precision")
+    try:
+      config.FLAGS.jax_default_dot_precision = "tensorfloat32"
+      jnp.dot(x, x)  # doesn't crash
+      jaxpr = jax.make_jaxpr(jnp.dot)(x, x)
+    finally:
+      config.FLAGS.jax_default_dot_precision = prev_val
+    self.assertIn('precision=HIGH', str(jaxpr))
+    self.assertEqual(prev_val, config._read("jax_default_dot_precision"))
+
+    prev_val = config._read("jax_default_dot_precision")
+    try:
+      config.update('jax_default_dot_precision','tensorfloat32')
+      jnp.dot(x, x)  # doesn't crash
+      jaxpr = jax.make_jaxpr(jnp.dot)(x, x)
+    finally:
+      config.update('jax_default_dot_precision', prev_val)
+    self.assertIn('precision=HIGH', str(jaxpr))
+    self.assertEqual(prev_val, config._read("jax_default_dot_precision"))
+
 
 class RematTest(jtu.JaxTestCase):