diff --git a/tensorflow_probability/python/bijectors/bijector.py b/tensorflow_probability/python/bijectors/bijector.py index d30183fcec..dd10b51ad5 100644 --- a/tensorflow_probability/python/bijectors/bijector.py +++ b/tensorflow_probability/python/bijectors/bijector.py @@ -44,6 +44,7 @@ ] JAX_MODE = False +NUMPY_MODE = False SKIP_DTYPE_CHECKS = False # Singleton object representing "no value", in cases where "None" is meaningful. @@ -1914,11 +1915,12 @@ def __str__(self): maybe_dtype = '' if self.forward_min_event_ndims == self.inverse_min_event_ndims: maybe_min_ndims = ', min_event_ndims={}'.format( - self.forward_min_event_ndims) + _unwrap_event_ndims(self.forward_min_event_ndims)) else: maybe_min_ndims = ( ', forward_min_event_ndims={}, inverse_min_event_ndims={}'.format( - self.forward_min_event_ndims, self.inverse_min_event_ndims)) + _unwrap_event_ndims(self.forward_min_event_ndims), + _unwrap_event_ndims(self.inverse_min_event_ndims))) maybe_min_ndims = maybe_min_ndims.replace('\'', '') return ('tfp.bijectors.{type_name}(' '"{self_name}"' @@ -1949,8 +1951,10 @@ def __repr__(self): type_name=type(self).__name__, self_name=self.name or '', batch_shape=batch_shape_str, - forward_min_event_ndims=self.forward_min_event_ndims, - inverse_min_event_ndims=self.inverse_min_event_ndims, + forward_min_event_ndims=_unwrap_event_ndims( + self.forward_min_event_ndims), + inverse_min_event_ndims=_unwrap_event_ndims( + self.inverse_min_event_ndims), dtype_x=_str_dtype(self.inverse_dtype()), dtype_y=_str_dtype(self.forward_dtype()))) @@ -2026,7 +2030,7 @@ def check_valid_ndims(ndims, validate=True): assertions = [] shape = ps.shape(ndims) - if not tf.is_tensor(shape): + if not tf.is_tensor(shape) or NUMPY_MODE or JAX_MODE: if shape.tolist(): raise ValueError('Expected scalar, saw shape {}.'.format(shape)) elif validate: @@ -2390,3 +2394,7 @@ def _str_tensorshape(x): if tensorshape_util.rank(x) is None: return '?' return str(tensorshape_util.as_list(x)).replace('None', '?') + + +def _unwrap_event_ndims(ndims): + return nest.map_structure(int, ndims) diff --git a/tensorflow_probability/python/bijectors/sinh_arcsinh.py b/tensorflow_probability/python/bijectors/sinh_arcsinh.py index 643c2be459..306307dcbc 100644 --- a/tensorflow_probability/python/bijectors/sinh_arcsinh.py +++ b/tensorflow_probability/python/bijectors/sinh_arcsinh.py @@ -99,6 +99,7 @@ def __init__(self, forward_min_event_ndims=0, validate_args=validate_args, parameters=parameters, + dtype=dtype, name=name) @classmethod diff --git a/tensorflow_probability/python/distributions/joint_distribution_coroutine_test.py b/tensorflow_probability/python/distributions/joint_distribution_coroutine_test.py index 93759cc814..35b2385509 100644 --- a/tensorflow_probability/python/distributions/joint_distribution_coroutine_test.py +++ b/tensorflow_probability/python/distributions/joint_distribution_coroutine_test.py @@ -989,7 +989,7 @@ def _get_support_bijectors(dists, xs=None, ys=None): event_ndims = [0, 1, 0] fldj = joint_bijector.forward_log_det_jacobian(xs, event_ndims) fldj_fd = _finite_difference_ldj(bijectors, 'forward', xs, delta=0.01) - self.assertAllClose(self.evaluate(fldj), self.evaluate(fldj_fd), rtol=1e-5) + self.assertAllClose(self.evaluate(fldj), self.evaluate(fldj_fd), rtol=2e-5) # Test inverse log det Jacobian via finite differences. ildj = joint_bijector.inverse_log_det_jacobian(ys, event_ndims) diff --git a/tensorflow_probability/python/distributions/markov_chain_test.py b/tensorflow_probability/python/distributions/markov_chain_test.py index a28b8a697c..35c8d9a263 100644 --- a/tensorflow_probability/python/distributions/markov_chain_test.py +++ b/tensorflow_probability/python/distributions/markov_chain_test.py @@ -106,7 +106,7 @@ def test_log_prob_matches_linear_gaussian_ssm(self): num_steps=7) x = markov_chain.sample(5, seed=seed) - self.assertAllClose(lgssm.log_prob(x), markov_chain.log_prob(x), rtol=1e-4) + self.assertAllClose(lgssm.log_prob(x), markov_chain.log_prob(x), rtol=5e-4) @test_util.numpy_disable_test_missing_functionality( 'JointDistributionNamedAutoBatched') diff --git a/tensorflow_probability/python/internal/backend/numpy/numpy_math.py b/tensorflow_probability/python/internal/backend/numpy/numpy_math.py index 2037d40e7e..7ecc67639f 100644 --- a/tensorflow_probability/python/internal/backend/numpy/numpy_math.py +++ b/tensorflow_probability/python/internal/backend/numpy/numpy_math.py @@ -300,7 +300,7 @@ def _invert_permutation(x, name=None): def _l2_normalize(x, axis=None, epsilon=1e-12, name=None): # pylint: disable=unused-argument x = _convert_to_tensor(x) norm = np.linalg.norm(x, ord=2, axis=_astuple(axis), keepdims=True) - norm = np.maximum(norm, np.sqrt(epsilon)) + norm = np.maximum(norm, np.sqrt(np.asarray(epsilon, dtype=norm.dtype))) return x / norm