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 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
autoray.register_function('tensorflow', 'linalg.qr', tf_qr)
import qubit_networks as beeky
import quimb as qu
import quimb.tensor as qtn
from quimb.tensor.optimize import TNOptimizer
LX, LY = 3, 2
DTYPE = 'complex128'
autodiff_backend = 'tensorflow'
autodiff_backend_opts = {'experimental_compile': True}
# autodiff_backend_opts = dict()
def state_energy(psi: beeky.QubitEncodeVector, hterms: dict, vterms: dict,
def _i(name):
"""Convenience function to import PennyLane
interfaces via a string pattern"""
if name == "tf":
return import_module("tensorflow")
if name == "qml":
return import_module("pennylane")
return import_module(name)
# -------------------------------- NumPy --------------------------------- #
from scipy.linalg import block_diag as _scipy_block_diag
ar.register_function("numpy", "flatten", lambda x: x.flatten())
ar.register_function("numpy", "coerce", lambda x: x)
ar.register_function("numpy", "block_diag", lambda x: _scipy_block_diag(*x))
ar.register_function("builtins", "block_diag", lambda x: _scipy_block_diag(*x))
ar.register_function("numpy", "gather", lambda x, indices: x[np.array(indices)])
ar.register_function("numpy", "unstack", list)
def _scatter_element_add_numpy(tensor, index, value):
"""In-place addition of a multidimensional value over various
indices of a tensor."""
new_tensor = tensor.copy()
new_tensor[tuple(index)] += value
return new_tensor