def testJacobian(self): bijector = tfb.MatrixInverseTriL() batch_size = 5 for ndims in range(2, 5): x_ = np.tril( np.random.uniform( -1., 1., size=[batch_size, ndims, ndims]).astype(np.float64)) fldj = bijector.forward_log_det_jacobian(x_, event_ndims=2) fldj_theoretical = bijector_test_util.get_fldj_theoretical( bijector, x_, event_ndims=2, input_to_unconstrained=tfb.Invert(tfb.FillTriangular()), output_to_unconstrained=tfb.Invert(tfb.FillTriangular())) fldj_, fldj_theoretical_ = self.evaluate([fldj, fldj_theoretical]) self.assertAllClose(fldj_, fldj_theoretical_)
def _forward(self, x): x = tf.convert_to_tensor(x) # Remove the first row and last column so we can extract strictly # lower triangular entries. n = x.shape[-1] t = tf.linalg.band_part(x[..., 1:, :-1], num_lower=n - 1, num_upper=0) return tfb.FillTriangular().inverse(t)
def testShape(self): x_shape = tf.TensorShape([5, 4, 6]) y_shape = tf.TensorShape([5, 4, 3, 3]) b = tfb.FillTriangular(validate_args=True) x = tf.ones(shape=x_shape, dtype=tf.float32) y_ = b.forward(x) self.assertAllEqual(tensorshape_util.as_list(y_.shape), tensorshape_util.as_list(y_shape)) x_ = b.inverse(y_) self.assertAllEqual(tensorshape_util.as_list(x_.shape), tensorshape_util.as_list(x_shape)) y_shape_ = b.forward_event_shape(x_shape) self.assertAllEqual(tensorshape_util.as_list(y_shape_), tensorshape_util.as_list(y_shape)) x_shape_ = b.inverse_event_shape(y_shape) self.assertAllEqual(tensorshape_util.as_list(x_shape_), tensorshape_util.as_list(x_shape)) y_shape_tensor = self.evaluate( b.forward_event_shape_tensor(tensorshape_util.as_list(x_shape))) self.assertAllEqual(y_shape_tensor, tensorshape_util.as_list(y_shape)) x_shape_tensor = self.evaluate( b.inverse_event_shape_tensor(tensorshape_util.as_list(y_shape))) self.assertAllEqual(x_shape_tensor, tensorshape_util.as_list(x_shape))
def testWithLKJSamples(self, dimension, concentration): bijector = tfb.CorrelationCholesky() lkj_dist = lkj.LKJ(dimension=dimension, concentration=np.float64(concentration), input_output_cholesky=True) batch_size = 10 y = self.evaluate(lkj_dist.sample([batch_size])) x = self.evaluate(bijector.inverse(y)) bijector_test_util.assert_bijective_and_finite(bijector, x, y, eval_func=self.evaluate, event_ndims=1, inverse_event_ndims=2, rtol=1e-5) fldj = bijector.forward_log_det_jacobian(x, event_ndims=1) fldj_theoretical = bijector_test_util.get_fldj_theoretical( bijector, x, event_ndims=1, inverse_event_ndims=2, output_to_unconstrained=tfb.Invert(tfb.FillTriangular())) self.assertAllClose(self.evaluate(fldj_theoretical), self.evaluate(fldj), atol=1e-5, rtol=1e-5)
def testShapeError(self): b = tfb.FillTriangular(validate_args=True) x_shape_bad = tf.TensorShape([5, 4, 7]) with self.assertRaisesRegexp(ValueError, "is not a triangular number"): b.forward_event_shape(x_shape_bad) with self.assertRaisesOpError("is not a triangular number"): self.evaluate(b.forward_event_shape_tensor(x_shape_bad.as_list())) y_shape_bad = tf.TensorShape([5, 4, 3, 2]) with self.assertRaisesRegexp(ValueError, "Matrix must be square"): b.inverse_event_shape(y_shape_bad) with self.assertRaisesOpError("Matrix must be square"): self.evaluate(b.inverse_event_shape_tensor(y_shape_bad.as_list()))
def testBijector(self): x = np.float32(np.array([1., 2., 3.])) y = np.float32(np.array([[3., 0.], [2., 1.]])) b = tfb.FillTriangular() y_ = self.evaluate(b.forward(x)) self.assertAllClose(y, y_) x_ = self.evaluate(b.inverse(y)) self.assertAllClose(x, x_) fldj = self.evaluate(b.forward_log_det_jacobian(x, event_ndims=1)) self.assertAllClose(fldj, 0.) ildj = self.evaluate(b.inverse_log_det_jacobian(y, event_ndims=2)) self.assertAllClose(ildj, 0.)
def testJacobian(self): cholesky_to_vector = tfb.Chain([ tfb.Invert(tfb.FillTriangular()), tfb.TransformDiagonal(tfb.Invert(tfb.Exp())) ]) bijector = tfb.CholeskyToInvCholesky() for x in [np.array([[2.]], dtype=np.float64), np.array([[2., 0.], [3., 4.]], dtype=np.float64), np.array([[2., 0., 0.], [3., 4., 0.], [5., 6., 7.]], dtype=np.float64)]: fldj = bijector.forward_log_det_jacobian(x, event_ndims=2) fldj_numerical = self._get_fldj_numerical( bijector, x, event_ndims=2, eps=1.e-6, input_to_vector=cholesky_to_vector, output_to_vector=cholesky_to_vector) fldj_, fldj_numerical_ = self.evaluate([fldj, fldj_numerical]) self.assertAllClose(fldj_, fldj_numerical_)
def testTheoreticalFldj(self): bijector = tfb.CorrelationCholesky() x = np.linspace(-50, 50, num=30).reshape(5, 6).astype(np.float64) y = self.evaluate(bijector.forward(x)) bijector_test_util.assert_bijective_and_finite(bijector, x, y, eval_func=self.evaluate, event_ndims=1, inverse_event_ndims=2, rtol=1e-5) fldj = bijector.forward_log_det_jacobian(x, event_ndims=1) fldj_theoretical = bijector_test_util.get_fldj_theoretical( bijector, x, event_ndims=1, inverse_event_ndims=2, output_to_unconstrained=tfb.Invert(tfb.FillTriangular())) self.assertAllClose(self.evaluate(fldj_theoretical), self.evaluate(fldj), atol=1e-5, rtol=1e-5)
def build_trainable_highway_flow(width, residual_fraction_initial_value=0.5, activation_fn=None, gate_first_n=None, seed=None, validate_args=False): """Builds a HighwayFlow parameterized by trainable variables. The variables are transformed to enforce the following parameter constraints: - `residual_fraction` is bounded between 0 and 1. - `upper_diagonal_weights_matrix` is a randomly initialized (lower) diagonal matrix with positive diagonal of size `width x width`. - `lower_diagonal_weights_matrix` is a randomly initialized lower diagonal matrix with ones on the diagonal of size `width x width`; - `bias` is a randomly initialized vector of size `width`. Args: width: Input dimension of the bijector. residual_fraction_initial_value: Initial value for gating parameter, must be between 0 and 1. activation_fn: Callable invertible activation function (e.g., `tf.nn.softplus`), or `None`. gate_first_n: Decides which part of the input should be gated (useful for example when using auxiliary variables). seed: Seed for random initialization of the weights. validate_args: Python `bool`. Whether to validate input with runtime assertions. Default value: `False`. Returns: trainable_highway_flow: The initialized bijector. """ residual_fraction_initial_value = tf.convert_to_tensor( residual_fraction_initial_value, dtype_hint=tf.float32, name='residual_fraction_initial_value') dtype = residual_fraction_initial_value.dtype bias_seed, upper_seed, lower_seed = samplers.split_seed(seed, n=3) lower_bijector = tfb.Chain([ tfb.TransformDiagonal(diag_bijector=tfb.Shift(1.)), tfb.Pad(paddings=[(1, 0), (0, 1)]), tfb.FillTriangular() ]) unconstrained_lower_initial_values = samplers.normal( shape=lower_bijector.inverse_event_shape([width, width]), mean=0., stddev=.01, seed=lower_seed) upper_bijector = tfb.FillScaleTriL(diag_bijector=tfb.Softplus(), diag_shift=None) unconstrained_upper_initial_values = samplers.normal( shape=upper_bijector.inverse_event_shape([width, width]), mean=0., stddev=.01, seed=upper_seed) return HighwayFlow(residual_fraction=util.TransformedVariable( initial_value=residual_fraction_initial_value, bijector=tfb.Sigmoid(), dtype=dtype), activation_fn=activation_fn, bias=tf.Variable(samplers.normal((width, ), mean=0., stddev=0.01, seed=bias_seed), dtype=dtype), upper_diagonal_weights_matrix=util.TransformedVariable( initial_value=upper_bijector.forward( unconstrained_upper_initial_values), bijector=upper_bijector, dtype=dtype), lower_diagonal_weights_matrix=util.TransformedVariable( initial_value=lower_bijector.forward( unconstrained_lower_initial_values), bijector=lower_bijector, dtype=dtype), gate_first_n=gate_first_n, validate_args=validate_args)