def test_qr_thin_square_fat(backend, shape): if backend == 'sparse': pytest.xfail("Sparse doesn't support linear algebra yet...") x = gen_rand(shape, backend) Q, R = ar.do('linalg.qr', x) xn, Qn, Rn = map(ar.to_numpy, (x, Q, R)) assert ar.do('allclose', xn, Qn @ Rn)
def _get_exp_stab_gate(self, coo, tau): ''' Returns ------- `exp(tau * multiplier * stabilizer)`: qarray Expm() of stabilizer centered on empty face at `coo`. Params ------- `coo`: tuple (x,y) (Irrelevant) location of the empty face that the stabilizer corresponds to. Just a label. `tau`: float Imaginary time for the exp(tau * gate) ''' key = (coo, tau) if key not in self._exp_stab_gates: where, gate = self._stab_gates[coo] el, ev = do('linalg.eigh', gate) expgate = ev @ do('diag', do('exp', el * tau)) @ dag(ev) self._exp_stab_gates[key] = (where, expgate) return self._exp_stab_gates[key]
def _expm_cached(self, x, y): cache = self._op_cache['expm'] key = (id(x), y) if key not in cache: el, ev = do('linalg.eigh', x) cache[key] = ev @ do('diag', do('exp', el * y)) @ dag(ev) return cache[key]
def norm_fro(x): if isinstance(x, numpy.ndarray): return norm_fro_dense(x.reshape(-1)) try: return do('linalg.norm', reshape(x, [-1]), 2) except AttributeError: return do('sum', do('multiply', do('conj', x), x)) ** 0.5
def test_linalg_svd_square(backend): if backend == 'sparse': pytest.xfail("Sparse doesn't support linear algebra yet...") x = gen_rand((5, 4), backend) U, s, V = ar.do('linalg.svd', x) assert (ar.infer_backend(x) == ar.infer_backend(U) == ar.infer_backend(s) == ar.infer_backend(V) == backend) y = U @ ar.do('diag', s, like=x) @ V diff = ar.do('sum', abs(y - x)) assert ar.to_numpy(diff) < 1e-8
def fsim_param_gen(params): theta, phi = params a = do('cos', theta) b = -1j * do('sin', theta) c = do('exp', -1j * phi) data = [[[[1, 0], [0, 0]], [[0, a], [b, 0]]], [[[0, b], [a, 0]], [[0, 0], [0, c]]]] return do('array', data, like=params)
def get_expm_gate(self, edge, t): '''Local term for `edge`, matrix-exponentiated by `t`. ''' # return self._expm_cached(self.get_gate(edge), x) key = (edge, t) if key not in self._exp_gates: gate = self.get_gate(edge) el, ev = do('linalg.eigh', gate) self._exp_gates[key] = ev @ do('diag', do('exp', el * t)) @ dag(ev) return self._exp_gates[key]
def modified_gram_schmidt(X): Q = [] for j in range(0, X.shape[0]): q = X[j, :] for i in range(0, j): rij = ar.do('tensordot', ar.do('conj', Q[i]), q, 1) q = q - rij * Q[i] rjj = ar.do('linalg.norm', q, 2) Q.append(q / rjj) return ar.do('stack', Q, axis=0, like=X)
def test_count_nonzero(backend, array_dtype): if backend == 'mars': import mars if mars._version.version_info < (0, 4, 0, ''): pytest.xfail('mars count_nonzero bug fixed in version 0.4.') if array_dtype == 'int': x = ar.do('array', [0, 1, 2, 0, 3], like=backend) elif array_dtype == 'float': x = ar.do('array', [0., 1., 2., 0., 3.], like=backend) elif array_dtype == 'bool': x = ar.do('array', [False, True, True, False, True], like=backend) nz = ar.do('count_nonzero', x) assert ar.to_numpy(nz) == 3
def _unitize_exp(x): r"""Perform isometrization using the using anti-symmetric matrix exponentiation. .. math:: U_A = \exp{A - A^\dagger} If ``x`` is rectangular it is completed with zeros first. """ m, n = x.shape d = max(m, n) x = do('pad', x, [[0, d - m], [0, d - n]], 'constant', constant_values=0.0) expx = do('linalg.expm', x - dag(x)) return expx[:m, :n]
def test_translator_random_normal(backend): from autoray import numpy as anp x = anp.random.normal(100.0, 0.1, size=(4, 5), like=backend) if backend == 'sparse': assert (x.data > 90.0).all() assert (x.data < 110.0).all() return assert (ar.to_numpy(x) > 90.0).all() assert (ar.to_numpy(x) < 110.0).all() if backend == 'tensorflow': x32 = ar.do('random.normal', 100.0, 0.1, dtype='float32', size=(4, 5), like=backend) assert x32.dtype == 'float32' assert (ar.to_numpy(x32) > 90.0).all() assert (ar.to_numpy(x32) < 110.0).all() # test default single scalar x = anp.random.normal(loc=1500, scale=10, like=backend) assert 1000 <= ar.to_numpy(x) < 2000
def recursively_stack_chunks(loc, rem): if not rem: return chunks[loc] return do('stack', [recursively_stack_chunks(loc + (d,), rem[1:]) for d in range(self.size_dict[rem[0]])], axis=output_pos[rem[0]] - len(loc), like=s)
def test_mgs(backend): if backend == 'sparse': pytest.xfail("Sparse doesn't support linear algebra yet...") x = gen_rand((3, 5), backend) Ux = modified_gram_schmidt(x) y = ar.do('sum', Ux @ ar.dag(Ux)) assert ar.to_numpy(y) == pytest.approx(3)
def state_energy(psi, hterms, vterms, **opts): he = psi.compute_local_expectation(hterms, normalized=True, **opts) ve = psi.compute_local_expectation(vterms, normalized=True, **opts) return autoray.do('real', (he + ve))
def test_complex_creation(backend, real_dtype): if backend == 'torch': pytest.xfail("Pytorch doesn't support complex numbers yet...") if (backend == 'sparse') and (real_dtype == 'float32'): pytest.xfail("Bug in sparse where single precision isn't maintained " "after scalar multiplication.") x = ar.do( 'complex', ar.astype(ar.do('random.normal', size=(3, 4), like=backend), real_dtype), ar.astype(ar.do('random.normal', size=(3, 4), like=backend), real_dtype)) assert ar.get_dtype_name(x) == { 'float32': 'complex64', 'float64': 'complex128' }[real_dtype]
def tf_qr(x): U, s, VH = autoray.do('linalg.svd', x) dtype = autoray.get_dtype_name(U) if 'complex' in dtype: s = autoray.astype(s, dtype) Q = U R = autoray.reshape(s, (-1, 1)) * VH return Q, R
def test_register_function(backend): x = ar.do('ones', shape=(2, 3), like=backend) def direct_fn(x): return 1 # first test we can provide the function directly ar.register_function(backend, 'test_register', direct_fn) assert ar.do('test_register', x) == 1 def wrap_fn(fn): def new_fn(*args, **kwargs): res = fn(*args, **kwargs) return res + 1 return new_fn # then check we can wrap the old (previous) function ar.register_function(backend, 'test_register', wrap_fn, wrap=True) assert ar.do('test_register', x) == 2
def test_dtype_specials(backend, creation, dtype): import numpy as np x = ar.do(creation, shape=(2, 3), like=backend) if backend == 'torch' and 'complex' in dtype: pytest.xfail("Pytorch doesn't support complex numbers yet...") x = ar.astype(x, dtype) assert ar.get_dtype_name(x) == dtype x = ar.to_numpy(x) assert isinstance(x, np.ndarray) assert ar.get_dtype_name(x) == dtype
def _unitize_qr(x): """Perform isometrization using the QR decomposition. """ fat = x.shape[0] < x.shape[1] if fat: x = transpose(x) Q = do('linalg.qr', x)[0] if fat: Q = transpose(Q) return Q
def state_energy(psi: beeky.QubitEncodeVector, hterms: dict, vterms: dict, **opts): '''Energy <psi|H|psi>, summing contribns from horiz/vertical terms in Hamiltonian. ''' # TODO: compute row/col envs first? he = psi.compute_local_expectation(hterms, normalized=True, **opts) ve = psi.compute_local_expectation(vterms, normalized=True, **opts) return autoray.do('real', (he + ve))
def test_triu(backend): x = gen_rand((4, 4), backend) xl = ar.do('triu', x) xln = ar.to_numpy(xl) assert xln[1, 0] == 0.0 if backend != 'sparse': # this won't work for sparse because density < 1 assert (xln > 0.0).sum() == 10 xl = ar.do('triu', x, k=-1) xln = ar.to_numpy(xl) if backend != 'sparse': # this won't work for sparse because density < 1 assert xln[1, 0] != 0.0 assert xln[2, 0] == 0.0 if backend != 'sparse': # this won't work for sparse because density < 1 assert (xln > 0.0).sum() == 13 if backend == 'tensorflow': with pytest.raises(ValueError): ar.do('triu', x, 1)
def _unitize_modified_gram_schmidt(A): """Perform isometrization explicitly using the modified Gram Schmidt procedure. """ m, n = A.shape thin = m > n if thin: A = do('transpose', A) Q = [] for j in range(0, min(m, n)): q = A[j, :] for i in range(0, j): rij = do('tensordot', do('conj', Q[i]), q, 1) q = q - rij * Q[i] Q.append(q / do('linalg.norm', q, 2)) Q = do('stack', Q, axis=0, like=A) if thin: Q = do('transpose', Q) return Q
def gen_rand(shape, backend, dtype='float64'): if backend == 'jax': from jax import random as jrandom global JAX_RANDOM_KEY if JAX_RANDOM_KEY is None: JAX_RANDOM_KEY = jrandom.PRNGKey(42) JAX_RANDOM_KEY, subkey = jrandom.split(JAX_RANDOM_KEY) return jrandom.uniform(subkey, shape=shape, dtype=dtype) elif backend == 'sparse': return ar.do('random.uniform', size=shape, like=backend, density=0.5, format='coo', fill_value=0) x = ar.do('random.uniform', size=shape, like=backend) x = ar.astype(x, ar.to_backend_dtype(dtype, backend)) assert ar.get_dtype_name(x) == dtype return x
def u3_gate_param_gen(params): theta, phi, lamda = params c2 = do('cos', theta / 2) s2 = do('sin', theta / 2) el = do('exp', 1.j * lamda) ep = do('exp', 1.j * phi) elp = do('exp', 1.j * (lamda + phi)) data = [[c2, -el * s2], [ep * s2, elp * c2]] return do('array', data, like=params)
def _svd(x, cutoff=-1.0, cutoff_mode=3, max_bond=-1, absorb=0): if isinstance(x, np.ndarray): return _svd_numpy(x, cutoff, cutoff_mode, max_bond, absorb) U, s, VH = do('linalg.svd', x) if cutoff > 0.0: if cutoff_mode == 1: n_chi = do('count_nonzero', s > cutoff) elif cutoff_mode == 2: n_chi = do('count_nonzero', s > cutoff * s[0]) elif cutoff_mode in (3, 4): s2 = s * s cs2 = do('cumsum', s2, 0) tot = cs2[-1] if cutoff_mode == 3: n_chi = do('count_nonzero', (tot - cs2) > cutoff) + 1 else: n_chi = do('count_nonzero', cs2 < (1 - cutoff) * tot) + 1 n_chi = max(n_chi, 1) if max_bond > 0: n_chi = min(n_chi, max_bond) if n_chi < s.shape[0]: s = s[:n_chi] U = U[..., :n_chi] VH = VH[:n_chi, ...] if cutoff_mode in (3, 4): norm = (tot / cs2[n_chi - 1])**0.5 s *= norm elif max_bond > 0: s = s[:max_bond] U = U[..., :max_bond] VH = VH[:max_bond, ...] if absorb == -1: U = U * reshape(s, (1, -1)) elif absorb == 1: VH = VH * reshape(s, (-1, 1)) else: s **= 0.5 U = U * reshape(s, (1, -1)) VH = VH * reshape(s, (-1, 1)) return U, VH
def ry_gate_param_gen(params): phi, = params c = do('cos', phi / 2) s = do('sin', phi / 2) data = [[c, -s], [s, c]] return do('array', data, like=params)
def constant(x, backend=_DEFAULT_BACKEND): if backend == 'jax' and isinstance(x, qarray): x = x.A return do('array', x, like=backend)
def _lq(x): if isinstance(x, np.ndarray): return _lq_numba(x) Q, L = do('linalg.qr', do('transpose', x)) return do('transpose', L), do('transpose', Q)
def _qr(x): if isinstance(x, np.ndarray): return _qr_numba(x) return do('linalg.qr', x)
def rz_gate_param_gen(params): phi, = params c = do('cos', phi / 2) s = -1j * do('sin', phi / 2) data = [[c + s, 0], [0, c - s]] return do('array', data, like=params)