Merge pull request #1174 from jburnim/r0.12

Prepare branch for the TFP 0.12.0rc2 release

Author: jburnim

STYLE_GUIDE.md

@@ -109,7 +109,7 @@ they supersede all previous conventions.
     *   Definitely use named args for 2nd args onward in docstrings.
 
 1.  Use names which describe semantics, not computation or mathematics, e.g.,
-    avoid `xp1 = x+1` or `tfd.Normal(loc=mu, scale=sigma)`.
+    avoid `xp1 = x + 1` or `tfd.Normal(loc=mu, scale=sigma)`.
 
 1.  Prefer inlining intermediates which are used once.
 
@@ -157,16 +157,16 @@ they supersede all previous conventions.
 
 1.  Prefer using the most specific TF operator. E.g,
 
-    *   Use `tf.squared_difference(x,y)` over `(x-y)**2`.
-    *   Use `tf.rsqrt` over `1./tf.sqrt(x)`.
+    *   Use `tf.squared_difference(x, y)` over `(x - y)**2`.
+    *   Use `tf.rsqrt` over `1. / tf.sqrt(x)`.
 
 1.  Worry about gradients! (It's often not automatic for API builders!)
 
 1.  When forced to choose between FLOPS and numerical accuracy, prefer numerical
     accuracy.
 
-1.  Avoid tf.cast if possible. Eg, prefer `tf.where(cond, a, b)` over
-    `tf.cast(cond,dtype=a.dtype)*a + (1-tf.cast(cond,dtype=b.dtype)*b`
+1.  Avoid tf.cast if possible. Eg, prefer `tf.where(pred, a, b)` over
+    `tf.cast(cond, dtype=a.dtype) * a + (1 - tf.cast(cond, dtype=b.dtype) * b`
 
 1.  Preserve static shape hints.
 
@@ -217,3 +217,15 @@ they supersede all previous conventions.
     `Tensor`s, and Numpy objects. When converting a user-provided literal to a
     `Tensor` (see e.g. `Distribution._call_log_prob`), specify the dtype to
     `tf.convert_to_tensor` if it is available.
+
+1.  Prefer overloaded operators on `Tensor`s (`+`, `-`, etc.) to explicit
+    method calls (`tf.add`, `tf.sub`, etc.). Exceptions:
+
+    * Prefer `tf.equal` to `==` when checking element-wise equality, because the
+      semantics of the latter are inconsistent between eager and graph
+      (`tf.function`) modes.
+    * Use `&` and `|` only if you want bitwise logic. Note that these are
+      equivalent to logical ops only if all inputs are `bool`s or are in
+      `{0, 1}`.
+
+

spinoffs/oryx/oryx/bijectors/__init__.py

@@ -18,24 +18,15 @@
 from oryx.bijectors import bijector_extensions
 from tensorflow_probability.substrates import jax as tfp
 
-__all__ = [
-    'bijector_extensions'
-]
-
 tfb = tfp.bijectors
 
-_bijectors = {}
+__all__ = tfb.__all__
 
-for name in dir(tfb):
+for name in __all__:
   bij = getattr(tfb, name)
   if inspect.isclass(bij) and issubclass(bij, tfb.Bijector):
     if bij is not tfb.Bijector:
       bij = bijector_extensions.make_type(bij)
-  _bijectors[name] = bij
-
-
-for key, val in _bijectors.items():
-  locals()[key] = val
-
+  locals()[name] = bij
 
-del _bijectors
+del tfb

spinoffs/oryx/oryx/core/interpreters/harvest.py

@@ -335,8 +335,8 @@ def process_higher_order_primitive(self, primitive, f, tracers, params,
       params = params.copy()
       new_params = dict(
           params,
-          mapped_invars=(True,) * len(tree_util.tree_leaves(plants)) +
-          params['mapped_invars'])
+          in_axes=(0,) * len(tree_util.tree_leaves(plants)) +
+          params['in_axes'])
     else:
       new_params = dict(params)
     all_args, all_tree = tree_util.tree_flatten((plants, vals))

spinoffs/oryx/oryx/core/interpreters/inverse/core.py

@@ -373,8 +373,8 @@ def remove_slice(cell):
   flat_vals, in_tree = tree_util.tree_flatten((mapped_incells, mapped_outcells))
   f, aux = flat_propagate(f, in_tree)
   # Assume all invars as mapped
-  new_mapped_invars = (True,) * len(flat_vals)
-  new_params = dict(params, mapped_invars=new_mapped_invars)
+  new_in_axes = (0,) * len(flat_vals)
+  new_params = dict(params, in_axes=new_in_axes)
   if 'donated_invars' in params:
     new_params['donated_invars'] = (False,) * len(flat_vals)
   subenv_vals = prim.bind(f, *flat_vals, **new_params)

spinoffs/oryx/oryx/core/interpreters/unzip.py

@@ -34,9 +34,9 @@
 from jax import core as jax_core
 from jax import custom_derivatives as cd
 from jax import linear_util as lu
-from jax import source_info_util
 from jax import tree_util
 from jax import util as jax_util
+from jax._src import source_info_util
 from jax.interpreters import partial_eval as pe
 import numpy as onp
 
@@ -282,14 +282,13 @@ def handle_call_primitive(self, call_primitive, f, tracers, params, is_map):
       return current_custom_rules()[call_primitive](self, f, *tracers, **params)
     if call_primitive in pe.call_partial_eval_rules:
       raise NotImplementedError
-    in_pvs, in_consts = jax_util.unzip2(t.pval for t in tracers)
+    in_pvals = [t.pval for t in tracers]
     if is_map:
-      pvs = [
-          None if pv is None else mapped_aval(params['axis_size'], pv)
-          for pv in in_pvs
-      ]
-    else:
-      pvs = in_pvs
+      unknown = pe.PartialVal.unknown
+      in_pvals = [pval if pval.is_known() or in_axis is None else
+                  unknown(mapped_aval(params['axis_size'], in_axis, pval[0]))
+                  for pval, in_axis in zip(in_pvals, params['in_axes'])]
+    pvs, in_consts = jax_util.unzip2(t.pval for t in tracers)
     keys = tuple(t.is_key() for t in tracers)
     new_settings = UnzipSettings(settings.tag, call_primitive in block_registry)
     fun, aux = unzip_eval(f, self, keys, tuple(pvs), new_settings)
@@ -360,12 +359,6 @@ def _bound_output_tracers(self, primitive, params, jaxpr, consts, env,
         for pv, const, key in safe_zip(out_pvs, out_consts, out_keys)
     ]
     new_params = dict(params, name=name, call_jaxpr=lifted_jaxpr)
-    if is_map:
-      new_params = dict(
-          new_params,
-          mapped_invars=tuple([True] * len(const_tracers) +
-                              [False] * len(env_tracers) +
-                              [True] * len(in_tracers)))
     if 'donated_invars' in params:
       new_donated_invars = (
           (False,) * len(const_tracers) + (False,) * len(env_tracers) +

spinoffs/oryx/oryx/distributions/__init__.py

@@ -16,23 +16,12 @@
 from oryx.distributions import distribution_extensions
 from tensorflow_probability.substrates import jax as tfp
 
-__all__ = [
-    'distribution_extensions'
-]
-
-
 tfd = tfp.distributions
 
-_distributions = {}
+__all__ = tfd.__all__
 
-for name in dir(tfd):
+for name in __all__:
   dist = getattr(tfd, name)
-  _distributions[name] = dist
-
-
-for key, val in _distributions.items():
-  locals()[key] = val
-
+  locals()[name] = dist
 
-del _distributions
-del distribution_extensions  # Only needed for registration.
+del tfd

spinoffs/oryx/oryx/experimental/nn/normalization_test.py

@@ -171,7 +171,7 @@ def test_check_grads(self):
     net = net_init.init(net_rng, state.Shape(in_shape))
 
     x = random.normal(data_rng, in_shape)
-    jtu.check_grads(net, (x,), 2)
+    jtu.check_grads(net.call, (x,), 2)
 
 
 def mse(x, y):

tensorflow_probability/python/__init__.py

@@ -18,20 +18,23 @@
 from __future__ import division
 from __future__ import print_function
 
+import functools
+
 from tensorflow_probability.python.internal import all_util
 from tensorflow_probability.python.internal import lazy_loader
 
 
-# Ensure TensorFlow is importable and its version is sufficiently recent. This
-# needs to happen before anything else, since the imports below will try to
-# import tensorflow, too.
 # pylint: disable=g-import-not-at-top
-def _ensure_tf_install():
-  """Attempt to import tensorflow, and ensure its version is sufficient.
+def _validate_tf_environment(package):
+  """Check TF version and (depending on package) warn about TensorFloat32.
+
+  Args:
+    package: Python `str` indicating which package is being imported. Used for
+      package-dependent warning about TensorFloat32.
 
   Raises:
     ImportError: if either tensorflow is not importable or its version is
-    inadequate.
+      inadequate.
   """
   try:
     import tensorflow.compat.v1 as tf
@@ -62,9 +65,10 @@ def _ensure_tf_install():
             required=required_tensorflow_version,
             present=tf.__version__))
 
-  if tf.config.experimental.tensor_float_32_execution_enabled():
+  if (package == 'mcmc' and
+      tf.config.experimental.tensor_float_32_execution_enabled()):
     # Must import here, because symbols get pruned to __all__.
-    import warnings  # pylint: disable=g-import-not-at-top
+    import warnings
     warnings.warn(
         'TensorFloat-32 matmul/conv are enabled for NVIDIA Ampere+ GPUs. The '
         'resulting loss of precision may hinder MCMC convergence. To turn off, '
@@ -94,6 +98,8 @@ def _ensure_tf_install():
 for pkg in _allowed_symbols:
   globals()[pkg] = lazy_loader.LazyLoader(
       pkg, globals(), 'tensorflow_probability.python.{}'.format(pkg),
-      on_first_access=_ensure_tf_install)
+      # These checks need to happen before lazy-loading, since the modules
+      # themselves will try to import tensorflow, too.
+      on_first_access=functools.partial(_validate_tf_environment, pkg))
 
 all_util.remove_undocumented(__name__, _allowed_symbols)

tensorflow_probability/python/bijectors/__init__.py

@@ -40,7 +40,6 @@
 from tensorflow_probability.python.bijectors.expm1 import Log1p
 from tensorflow_probability.python.bijectors.ffjord import FFJORD
 from tensorflow_probability.python.bijectors.fill_scale_tril import FillScaleTriL
-from tensorflow_probability.python.bijectors.fill_scale_tril import ScaleTriL
 from tensorflow_probability.python.bijectors.fill_triangular import FillTriangular
 from tensorflow_probability.python.bijectors.frechet_cdf import FrechetCDF
 from tensorflow_probability.python.bijectors.generalized_pareto import GeneralizedPareto
@@ -159,7 +158,6 @@
     "ScaleMatvecLinearOperatorBlock",
     "ScaleMatvecLU",
     "ScaleMatvecTriL",
-    "ScaleTriL",
     "Shift",
     "ShiftedGompertzCDF",
     "Sigmoid",

tensorflow_probability/python/bijectors/bijector_properties_test.py

@@ -77,7 +77,6 @@
     'ScaleMatvecTriL',
     'Shift',
     'ShiftedGompertzCDF',
-    'ScaleTriL',
     'Sigmoid',
     'Sinh',
     'SinhArcsinh',