def _to_list(x): if B.rank(x) == 0: return [x] elif B.rank(x) == 1: return list(x) else: raise ValueError(f'Could not convert "{x}" to a list.')
def convert(x): if B.rank(x) == 0: if isinstance(x, B.Number) and x == 0: return Zero(B.dtype(x), 1, 1) else: return Constant(x, 1, 1) elif B.rank(x) == 2: return Dense(x) else: raise RuntimeError( f"Cannot convert rank {B.rank(x)} input to a matrix.")
def assert_matrix(x, message): """Assert that a tensor is a matrix. Args x (tensor): Tensor that must be a matrix. message (str): Error message to raise in the case that `x` is not a matrix. """ if B.rank(x) != 2: raise AssertionError(message)
def assert_vector(x, message): """Assert that a tensor is a vector. Args x (tensor): Tensor that must be a vector. message (str): Error message to raise in the case that `x` is not a vector. """ if B.rank(x) != 1: raise AssertionError(message)
def wrapped_f(*args): converted_args = () for i, arg in enumerate(args): # Check if the argument is already a Pandas object. if isinstance(arg, (pd.Series, pd.DataFrame)): converted_args += (arg, ) continue # It is not already a Pandas object. Attempt to convert. if B.rank(arg) == 0: converted_args += (arg, ) elif B.rank(arg) == 1: converted_args += (pd.Series(arg), ) elif B.rank(arg) == 2: converted_args += (pd.DataFrame(arg), ) else: raise ValueError(f"Argument {i} has rank {B.rank(arg)}, which " f"cannot be automatically converted to a " f"Pandas object.") return f(*converted_args)
def eig(a, compute_eigvecs=True): if compute_eigvecs: vals, vecs = B.eig(a, compute_eigvecs=True) vals = B.flatten(vals) if B.rank(vecs) == 3: vecs = B.transpose(vecs, perm=(1, 0, 2)) vecs = B.reshape(vecs, 3, -1) order = compute_order(vals) return B.take(vals, order), B.abs(B.take(vecs, order, axis=1)) else: vals = B.flatten(B.eig(a, compute_eigvecs=False)) return B.take(vals, compute_order(vals))
def test_packer(): a, b, c = B.randn(5, 10), B.randn(20), B.randn(5, 1, 15) for packer, args in zip( [Packer(a, b, c), Packer([a, b, c])], [(a, b, c), ((a, b, c), )]): # Test packing. packed = packer.pack(*args) assert B.rank(packed) == 1 # Test unpacking. a_, b_, c_ = packer.unpack(packed) allclose(a, a_) allclose(b, b_) allclose(c, c_)
def test_ew_1d(check_lazy_shapes, batch_a, batch_b): a = Tensor(*((3,) if batch_a else ()), 10, 1).np() b = Tensor(*((3,) if batch_b else ()), 10, 1).np() # Check that we can feed both rank 1 and rank 2 tensors. for squeeze_a in [True, False]: for squeeze_b in [True, False]: if squeeze_a and B.rank(a) == 2: a2 = a[..., 0] else: a2 = a if squeeze_b and B.rank(b) == 2: b2 = b[..., 0] else: b2 = b approx(B.ew_dists2(a2, b2), np.abs(a - b) ** 2) approx(B.ew_dists2(a2), np.zeros((*((3,) if batch_a else ()), 10, 1))) approx(B.ew_dists(a2, b2), np.abs(a - b)) approx(B.ew_dists(a2), np.zeros((*((3,) if batch_a else ()), 10, 1))) approx(B.ew_sums2(a2, b2), np.abs(a + b) ** 2) approx(B.ew_sums2(a2), np.abs(a + a) ** 2) approx(B.ew_sums(a2, b2), np.abs(a + b)) approx(B.ew_sums(a2), np.abs(a + a))
def test_pack_unpack(): a, b, c = B.randn(5, 10), B.randn(20), B.randn(5, 1, 15) # Test packing. package = pack(a, b, c) assert B.rank(package) == 1 # Test unpacking. a2, b2, c2 = unpack(package, B.shape(a), B.shape(b), B.shape(c)) approx(a, a2) approx(b, b2) approx(c, c2) # Check that the package must be a vector. with pytest.raises(ValueError): unpack(B.randn(2, 2), (2, 2))
def hessian(f, x): """Compute the Hessian of a function at a certain input. Args: f (function): Function to compute Hessian of. x (column vector): Input to compute Hessian at. differentiable (bool, optional): Make the computation of the Hessian differentiable. Defaults to `False`. Returns: matrix: Hessian. """ if B.rank(x) != 2 or B.shape(x)[1] != 1: raise ValueError("Input must be a column vector.") # Use RMAD twice to preserve memory. hess = jax.jacrev(jax.jacrev(lambda x: f(x[:, None])))(x[:, 0]) return (hess + B.transpose(hess) ) / 2 # Symmetrise to counteract numerical errors.
def unpack(package: B.Numeric, *shapes): """Unpack vector. Args: package (tensor): Tensor to unpack. *shapes (shape): Shapes of objects to unpack. Returns: list[tensor]: Original objects. """ if B.rank(package) != 1: raise ValueError("Package must be a vector.") # Unpack package. lengths = [reduce(mul, shape, 1) for shape in shapes] i, outs = 0, [] for length, shape in zip(lengths, shapes): outs.append(B.reshape(package[i : i + length], *shape)) i += length return outs
def __call__(self, x): x_rank = B.rank(x) if x_rank == 1: x = x[None, :, None] elif x_rank == 2: x = x[None, :, :] elif x_rank == 3: pass else: raise ValueError("Cannot handle inputs of rank {}.".format(x_rank)) # Apply layers one by one. for layer in self.layers: x = layer(x) # Remove batch dimension, if that wasn't specified initially. if x_rank != 3: x = x[0, :, :] return x
def summarise_samples(x, samples, db=False): """Summarise samples. Args: x (vector): Inputs of samples. samples (tensor): Samples, with the first dimension corresponding to different samples. db (bool, optional): Convert to decibels. Returns: :class:`collections.namedtuple`: Named tuple containing various statistics of the samples. """ x, samples = B.to_numpy(x, samples) random_inds = np.random.permutation(B.shape(samples)[0])[:3] def transform(x): if db: return 10 * np.log10(x) else: return x perm = tuple(reversed(range(B.rank(samples)))) # Reverse all dimensions. return collect( x=B.to_numpy(x), mean=transform(B.mean(samples, axis=0)), var=transform(B.std(samples, axis=0))**2, err_68_lower=transform(B.quantile(samples, 0.32, axis=0)), err_68_upper=transform(B.quantile(samples, 1 - 0.32, axis=0)), err_95_lower=transform(B.quantile(samples, 0.025, axis=0)), err_95_upper=transform(B.quantile(samples, 1 - 0.025, axis=0)), err_99_lower=transform(B.quantile(samples, 0.0015, axis=0)), err_99_upper=transform(B.quantile(samples, 1 - 0.0015, axis=0)), samples=transform(B.transpose(samples, perm=perm)[..., random_inds]), all_samples=transform(B.transpose(samples, perm=perm)), )
def __call__(self, x): # Put the batch dimension second. x_rank = B.rank(x) if x_rank == 2: x = x[:, None, :] elif x_rank == 3: x = B.transpose(x, perm=(1, 0, 2)) else: raise ValueError(f"Cannot handle inputs of rank {B.rank(x)}.") # Recurrently apply the cell. n, batch_size, m = B.shape(x) y0 = B.zeros(B.dtype(x), batch_size, self.cell.width) h0 = B.tile(self.h0, batch_size, 1) res = B.scan(self.cell, x, h0, y0)[1] # Put the batch dimension first again. res = B.transpose(res, perm=(1, 0, 2)) # Remove the batch dimension, if that didn't exist before. if x_rank == 2: res = res[0, :, :] return res