def test_lowrank_attributes(): # Check default right and middle factor. left = B.ones(3, 2) lr = LowRank(left) assert lr.left is left assert lr.right is left approx(lr.middle, B.eye(2)) assert lr.rank == 2 # Check given identical right. lr = LowRank(left, left) assert lr.left is left assert lr.right is left approx(lr.middle, B.eye(2)) assert lr.rank == 2 # Check given identical right and middle. middle = B.ones(2, 2) lr = LowRank(left, left, middle=middle) assert lr.left is left assert lr.right is left approx(lr.middle, B.ones(2, 2)) assert lr.rank == 2 # Check given other right and middle factor. right = 2 * B.ones(3, 2) lr = LowRank(left, right, middle=middle) assert lr.left is left assert lr.right is right approx(lr.middle, B.ones(2, 2)) assert lr.rank == 2 # Check rank in non-square case. assert LowRank(B.ones(3, 2), B.ones(3, 2), B.ones(2, 2)).rank == 2 assert LowRank(B.ones(3, 2), B.ones(3, 1), B.ones(2, 1)).rank == 1
def test_normal_lazy_var_diag(): # If `var_diag` isn't set, the variance will be constructed to get the diagonal. dist = Normal(lambda: B.eye(3)) approx(dist.var_diag, B.ones(3)) approx(dist._var, B.eye(3)) # If `var_diag` is set, the variance will _not_ be constructed to get the diagonal. dist = Normal(lambda: B.eye(3), var_diag=lambda: 9) approx(dist.var_diag, 9) assert dist._var is None
def test_normal_lazy_nonzero_mean(): dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3)) # Nothing should be populated yet. assert dist._mean is None assert dist._var is None # But they should be populated upon request. approx(dist.mean, B.ones(3, 1)) assert dist._var is None approx(dist.var, B.eye(3))
def test_cholesky_retry_factor(check_lazy_shapes): # Try `cholesky_retry_factor = 1`. B.cholesky_retry_factor = 1 B.cholesky(B.zeros(3, 3)) B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * B.epsilon) with pytest.raises(np.linalg.LinAlgError): B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 10 * B.epsilon) with pytest.raises(np.linalg.LinAlgError): B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 100 * B.epsilon) # Try `cholesky_retry_factor = 10`. B.cholesky_retry_factor = 10 B.cholesky(B.zeros(3, 3)) B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * B.epsilon) B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 10 * B.epsilon) with pytest.raises(np.linalg.LinAlgError): B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 100 * B.epsilon) # Try `cholesky_retry_factor = 100`. B.cholesky_retry_factor = 100 B.cholesky(B.zeros(3, 3)) B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * B.epsilon) B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 10 * B.epsilon) B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 100 * B.epsilon) # Reset the factor! B.cholesky_retry_factor = 1
def test_normal_lazy_nonzero_mean(): dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3)) assert not dist.mean_is_zero approx(dist._mean, B.ones(3, 1)) assert dist._var is None approx(dist.mean, B.ones(3, 1)) assert dist._var is None approx(dist.var, B.eye(3))
def test_normal_mean_is_zero(): # Check zero case. dist = Normal(B.eye(3)) assert dist.mean_is_zero approx(dist.mean, B.zeros(3, 1)) # Check another zero case. dist = Normal(Zero(np.float32, 3, 1), B.eye(3)) assert dist.mean_is_zero approx(dist.mean, B.zeros(3, 1)) # Check nonzero case. assert not Normal(B.randn(3, 1), B.eye(3)).mean_is_zero
def test_normal_lazy_zero_mean(): dist = Normal(lambda: B.eye(3)) assert dist.mean_is_zero assert dist._mean is 0 assert dist._var is None approx(dist.mean, B.zeros(3, 1)) # At this point, the variance should be constructed, because it is used to get the # dimensionality and data type for the mean. assert dist._var is not None approx(dist.var, B.eye(3))
def assert_orthogonal(x): """Assert that a matrix is orthogonal.""" # Check that matrix is square. assert B.shape(x)[0] == B.shape(x)[1] # Check that its transpose is its inverse. approx(B.matmul(x, x, tr_a=True), B.eye(x))
def test_normal_sampling(): for mean in [0, 1]: dist = Normal(mean, 3 * B.eye(np.int32, 200)) # Sample without noise. samples = dist.sample(2000) approx(B.mean(samples), mean, atol=5e-2) approx(B.std(samples)**2, 3, atol=5e-2) # Sample with noise samples = dist.sample(2000, noise=2) approx(B.mean(samples), mean, atol=5e-2) approx(B.std(samples)**2, 5, atol=5e-2)
def test_cholesky_solve_ut(dense_pd): chol = B.cholesky(dense_pd) with AssertDenseWarning( [ "solving <upper-triangular> x = <diagonal>", "matrix-multiplying <upper-triangular> and <lower-triangular>", ] ): approx( B.cholesky_solve(B.transpose(chol), B.eye(chol)), B.inv(B.matmul(chol, chol, tr_a=True)), )
def pd_inv(a: Union[B.Numeric, AbstractMatrix]): """Invert a positive-definite matrix. Args: a (matrix): Positive-definite matrix to invert. Returns: matrix: Inverse of `a`, which is also positive definite. """ a = convert(a, AbstractMatrix) # The call to `cholesky_solve` will convert the identity matrix to dense, because # `cholesky(a)` will not have any exploitable structure. We suppress the expected # warning by converting `B.eye(a)` to dense here already. return B.cholesky_solve(B.cholesky(a), B.dense(B.eye(a)))
def test_normal_lazy_mean_var(): # The lazy `mean_var` should only be called when neither the mean nor the variance # exists. Otherwise, it's more efficient to just construct the other one. We # go over all branches in the `if`-statement. dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9)) approx(dist.mean_var, (8, 9)) approx(dist.mean, 8) approx(dist.var, 9) dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9)) approx(dist.mean, B.ones(3, 1)) approx(dist.mean_var, (B.ones(3, 1), B.eye(3))) approx(dist.var, B.eye(3)) dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9)) approx(dist.var, B.eye(3)) approx(dist.mean_var, (B.ones(3, 1), B.eye(3))) approx(dist.mean, B.ones(3, 1)) dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9)) approx(dist.var, B.eye(3)) approx(dist.mean, B.ones(3, 1)) approx(dist.mean_var, (B.ones(3, 1), B.eye(3)))
def test_normal_lazy_mean_var_diag(): # The lazy `mean_var_diag` should only be called when neither the mean nor the # diagonal of the variance exists. Otherwise, it's more efficient to just construct # the other one. We go over all branches in the `if`-statement. dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var_diag=lambda: (8, 9)) approx(dist.marginals(), (8, 9)) approx(dist.mean, 8) approx(dist.var_diag, 9) dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var_diag=lambda: (8, 9)) approx(dist.mean, B.ones(3, 1)) approx(dist.marginals(), (B.ones(3), B.ones(3))) approx(dist.var_diag, B.ones(3)) dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var_diag=lambda: (8, 9)) approx(dist.var_diag, B.ones(3)) approx(dist.marginals(), (B.ones(3), B.ones(3))) approx(dist.mean, B.ones(3, 1)) dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var_diag=lambda: (8, 9)) approx(dist.var_diag, B.ones(3)) approx(dist.mean, B.ones(3, 1)) approx(dist.marginals(), (B.ones(3), B.ones(3)))
def test_normal_sampling(): for mean in [0, 1]: dist = Normal(mean, 3 * B.eye(np.int32, 200)) # Sample without noise. samples = dist.sample(2000) approx(B.mean(samples), mean, atol=5e-2) approx(B.std(samples) ** 2, 3, atol=5e-2) # Sample with noise samples = dist.sample(2000, noise=2) approx(B.mean(samples), mean, atol=5e-2) approx(B.std(samples) ** 2, 5, atol=5e-2) state, sample1 = dist.sample(B.create_random_state(B.dtype(dist), seed=0)) state, sample2 = dist.sample(B.create_random_state(B.dtype(dist), seed=0)) assert isinstance(state, B.RandomState) approx(sample1, sample2)
def transform(x): tril = B.vec_to_tril(x, offset=-1) skew = tril - B.transpose(tril) eye = B.eye(skew) return B.solve(eye + skew, eye - skew)
def f1(x): dists2 = (x - B.transpose(x))**2 K = B.exp(-0.5 * dists2) K = K + B.epsilon * B.eye(t, n) L = B.cholesky(K) return B.matmul(L, B.ones(t, n, m))
def test_eye(dense1): approx(B.eye(dense1), np.eye(B.shape(dense1)[0])) assert isinstance(B.eye(dense1), Diagonal)
def test_logm(check_lazy_shapes): mat = B.eye(3) + 0.1 * B.randn(3, 3) check_sensitivity(logm, s_logm, (mat, )) check_grad(logm, (mat, ))
def test_cholesky_solve_lt(dense_pd): chol = B.cholesky(dense_pd) with AssertDenseWarning("solving <lower-triangular> x = <diagonal>"): approx(B.cholesky_solve(chol, B.eye(chol)), B.inv(dense_pd))
for a in Tensor(2, 2).forms(): assert "cpu" in str(B.device(a)).lower() approx(B.to_active_device(a), a) # Check that numbers remain unchanged. a = 1 assert B.to_active_device(a) is a @pytest.mark.parametrize("t", [tf.float32, torch.float32, jnp.float32]) @pytest.mark.parametrize( "f", [ lambda t: B.zeros(t, 2, 2), lambda t: B.ones(t, 2, 2), lambda t: B.eye(t, 2), lambda t: B.linspace(t, 0, 5, 10), lambda t: B.range(t, 10), lambda t: B.rand(t, 10), lambda t: B.randn(t, 10), ], ) def test_on_device(f, t, check_lazy_shapes): f_t = f(t) # Contruct on current and existing device. # Set the active device to something else. B.ActiveDevice.active_name = "previous" # Check that explicit allocation on CPU works. with B.on_device("cpu"): assert B.device(f(t)) == B.device(f_t)
def generate(code): """Generate a random tensor of a particular type, specified with a code. Args: code (str): Code of the matrix. Returns: tensor: Random tensor. """ mat_code, shape_code = code.split(":") # Parse shape. if shape_code == "": shape = () else: shape = tuple(int(d) for d in shape_code.split(",")) if mat_code == "randn": return B.randn(*shape) elif mat_code == "randn_pd": mat = B.randn(*shape) # If it is a scalar or vector, just pointwise square it. if len(shape) in {0, 1}: return mat**2 + 1 else: return B.matmul(mat, mat, tr_b=True) + B.eye(shape[0]) elif mat_code == "zero": return Zero(B.default_dtype, *shape) elif mat_code == "const": return Constant(B.randn(), *shape) elif mat_code == "const_pd": return Constant(B.randn()**2 + 1, *shape) elif mat_code == "lt": mat = B.vec_to_tril(B.randn(int(0.5 * shape[0] * (shape[0] + 1)))) return LowerTriangular(mat) elif mat_code == "lt_pd": mat = generate(f"randn_pd:{shape[0]},{shape[0]}") return LowerTriangular(B.cholesky(B.reg(mat))) elif mat_code == "ut": mat = B.vec_to_tril(B.randn(int(0.5 * shape[0] * (shape[0] + 1)))) return UpperTriangular(B.transpose(mat)) elif mat_code == "ut_pd": mat = generate(f"randn_pd:{shape[0]},{shape[0]}") return UpperTriangular(B.transpose(B.cholesky(B.reg(mat)))) elif mat_code == "dense": return Dense(generate(f"randn:{shape_code}")) elif mat_code == "dense_pd": return Dense(generate(f"randn_pd:{shape_code}")) elif mat_code == "diag": return Diagonal(generate(f"randn:{shape_code}")) elif mat_code == "diag_pd": return Diagonal(generate(f"randn_pd:{shape_code}")) else: raise RuntimeError(f'Cannot parse generation code "{code}".')
def test_logm(check_lazy_shapes): mat = B.eye(3) + 0.1 * B.randn(3, 3) check_function(B.logm, (Tensor(mat=mat),))
def inverse_transform(x): eye = B.eye(x) skew = B.solve(eye + x, eye - x) return B.tril_to_vec(skew, offset=-1)
def test_normal_dtype(normal1): assert B.dtype(Normal(0, B.eye(3))) == np.float64 assert B.dtype(Normal(B.ones(3), B.zeros(int, 3))) == np.float64 assert B.dtype(Normal(B.ones(int, 3), B.zeros(int, 3))) == np.int64