def optimize_compare(self, equation, operands=None, verbose=False):
for clean in [False, True]:
with self.subTest(equation=equation):
if operands is not None:
inputs = operands
else:
eqs = equation.split("->")[0].split(",")
inputs = []
for d, eq in enumerate(eqs):
i = numpy.arange(2**len(eq)).reshape(
(2, ) * len(eq)).astype(numpy.float32)
inputs.append(i +
numpy.array([3**d], dtype=numpy.float32))
exp = numpy.einsum(equation, *inputs)
if verbose:
print("###### equation", equation)
path = numpy.einsum_path(equation, *inputs, optimize=False)
print(path[1])
path = numpy.einsum_path(equation, *inputs)
print(path[1])
shapes = [m.shape for m in inputs]
vv = 12 if equation == ",a,ab,abc->abc" else verbose
with self.subTest(strategy='numpy'):
seq = decompose_einsum_equation(equation,
*shapes,
verbose=verbose,
strategy='numpy',
clean=clean)
got = apply_einsum_sequence(seq, *inputs, verbose=vv)
self.assertEqualArray(exp, got, decimal=6)
if clean:
with self.subTest(strategy='onnx'):
inps = ['X%d' % (i + 1) for i in range(len(inputs))]
try:
onx = seq.to_onnx('Y', *inps, dtype=numpy.float32)
except NotImplementedError as e:
if "diagonal" in str(e):
onx = None
else:
raise e
if onx is not None:
oinf = OnnxInference(onx)
inps = {
n: v.astype(numpy.float32)
for n, v in zip(inps, inputs)
}
got = oinf.run(inps, verbose=vv, fLOG=print)['Y']
self.assertEqualArray(exp, got, decimal=5)
with self.subTest(strategy='simple'):
seq = decompose_einsum_equation(equation,
*shapes,
clean=clean,
verbose=verbose)
got = apply_einsum_sequence(seq, *inputs, verbose=verbose)
self.assertEqualArray(exp, got, decimal=6)
def test_long_paths(self):
# Long complex cases
# Long test 1
long_test1 = self.build_operands(
'acdf,jbje,gihb,hfac,gfac,gifabc,hfac')
path, path_str = np.einsum_path(*long_test1, optimize='greedy')
self.assert_path_equal(
path,
['einsum_path', (3, 6), (3, 4), (2, 4), (2, 3), (0, 2), (0, 1)])
path, path_str = np.einsum_path(*long_test1, optimize='optimal')
self.assert_path_equal(
path,
['einsum_path', (3, 6), (3, 4), (2, 4), (2, 3), (0, 2), (0, 1)])
# Long test 2
long_test2 = self.build_operands('chd,bde,agbc,hiad,bdi,cgh,agdb')
path, path_str = np.einsum_path(*long_test2, optimize='greedy')
print(path)
self.assert_path_equal(
path,
['einsum_path', (3, 4), (0, 3), (3, 4), (1, 3), (1, 2), (0, 1)])
path, path_str = np.einsum_path(*long_test2, optimize='optimal')
print(path)
self.assert_path_equal(
path,
['einsum_path', (0, 5), (1, 4), (3, 4), (1, 3), (1, 2), (0, 1)])
def __call__(self, stage, itr):
"""!Record the spin-spin correlators."""
P = self.particle(stage, itr)
H = self.hole(stage, itr)
nx = P.shape[0]
nt = P.shape[1]
d = np.eye(nx*nt).reshape(*P.shape) # A Kronecker delta
log = getLogger(__name__)
data={}
data["np"] = d-P
data["nh"] = d-H
if self._einsum_path is None:
if self.transform is None:
# No need for the transformation, cut the cost:
self._einsum_path, _ = np.einsum_path("xtxt->x", data['np'], optimize="optimal")
log.info("Optimized Einsum path for time averaging.")
else:
# We'll time average and transform at once:
self._einsum_path, _ = np.einsum_path("ax,xtxt->a", self.transform, data["np"], optimize="optimal")
log.info("Optimized Einsum path for time averaging and unitary transformation.")
if self.transform is None:
for name, correlator in data.items():
measurement = np.einsum("xtxt->x", correlator, optimize=self._einsum_path) / nt
self.correlators[name].append(measurement)
else:
for name, correlator in data.items():
measurement = np.einsum("ax,xtxt->a", self.transform, correlator, optimize=self._einsum_path) / nt
self.correlators[name].append(measurement)
def test_edge_paths(self):
# Difficult edge cases
# Edge test1
edge_test1 = self.build_operands('eb,cb,fb->cef')
path, path_str = np.einsum_path(*edge_test1, optimize='greedy')
self.assert_path_equal(path, ['einsum_path', (0, 2), (0, 1)])
path, path_str = np.einsum_path(*edge_test1, optimize='optimal')
self.assert_path_equal(path, ['einsum_path', (0, 2), (0, 1)])
# Edge test2
edge_test2 = self.build_operands('dd,fb,be,cdb->cef')
path, path_str = np.einsum_path(*edge_test2, optimize='greedy')
self.assert_path_equal(path, ['einsum_path', (0, 3), (0, 1), (0, 1)])
path, path_str = np.einsum_path(*edge_test2, optimize='optimal')
self.assert_path_equal(path, ['einsum_path', (0, 3), (0, 1), (0, 1)])
# Edge test3
edge_test3 = self.build_operands('bca,cdb,dbf,afc->')
path, path_str = np.einsum_path(*edge_test3, optimize='greedy')
self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 2), (0, 1)])
path, path_str = np.einsum_path(*edge_test3, optimize='optimal')
self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 2), (0, 1)])
# Edge test4
edge_test4 = self.build_operands('dcc,fce,ea,dbf->ab')
path, path_str = np.einsum_path(*edge_test4, optimize='greedy')
self.assert_path_equal(path, ['einsum_path', (0, 3), (0, 2), (0, 1)])
path, path_str = np.einsum_path(*edge_test4, optimize='optimal')
self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 2), (0, 1)])
def test_edge_paths(self):
# Difficult edge cases
# Edge test1
edge_test1 = self.build_operands('eb,cb,fb->cef')
path, path_str = np.einsum_path(*edge_test1, optimize='greedy')
self.assert_path_equal(path, ['einsum_path', (0, 2), (0, 1)])
path, path_str = np.einsum_path(*edge_test1, optimize='optimal')
self.assert_path_equal(path, ['einsum_path', (0, 2), (0, 1)])
# Edge test2
edge_test2 = self.build_operands('dd,fb,be,cdb->cef')
path, path_str = np.einsum_path(*edge_test2, optimize='greedy')
self.assert_path_equal(path, ['einsum_path', (0, 3), (0, 1), (0, 1)])
path, path_str = np.einsum_path(*edge_test2, optimize='optimal')
self.assert_path_equal(path, ['einsum_path', (0, 3), (0, 1), (0, 1)])
# Edge test3
edge_test3 = self.build_operands('bca,cdb,dbf,afc->')
path, path_str = np.einsum_path(*edge_test3, optimize='greedy')
self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 2), (0, 1)])
path, path_str = np.einsum_path(*edge_test3, optimize='optimal')
self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 2), (0, 1)])
# Edge test4
edge_test4 = self.build_operands('dcc,fce,ea,dbf->ab')
path, path_str = np.einsum_path(*edge_test4, optimize='greedy')
self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 1), (0, 1)])
path, path_str = np.einsum_path(*edge_test4, optimize='optimal')
self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 2), (0, 1)])
# Edge test5
edge_test4 = self.build_operands('a,ac,ab,ad,cd,bd,bc->',
size_dict={
"a": 20,
"b": 20,
"c": 20,
"d": 20
})
path, path_str = np.einsum_path(*edge_test4, optimize='greedy')
self.assert_path_equal(path,
['einsum_path', (0, 1), (0, 1, 2, 3, 4, 5)])
path, path_str = np.einsum_path(*edge_test4, optimize='optimal')
self.assert_path_equal(path,
['einsum_path', (0, 1), (0, 1, 2, 3, 4, 5)])
def test_edge_paths(self):
# Difficult edge cases
# Edge test1
edge_test1 = self.build_operands("eb,cb,fb->cef")
path, path_str = np.einsum_path(*edge_test1, optimize="greedy")
self.assert_path_equal(path, ["einsum_path", (0, 2), (0, 1)])
path, path_str = np.einsum_path(*edge_test1, optimize="optimal")
self.assert_path_equal(path, ["einsum_path", (0, 2), (0, 1)])
# Edge test2
edge_test2 = self.build_operands("dd,fb,be,cdb->cef")
path, path_str = np.einsum_path(*edge_test2, optimize="greedy")
self.assert_path_equal(path, ["einsum_path", (0, 3), (0, 1), (0, 1)])
path, path_str = np.einsum_path(*edge_test2, optimize="optimal")
self.assert_path_equal(path, ["einsum_path", (0, 3), (0, 1), (0, 1)])
# Edge test3
edge_test3 = self.build_operands("bca,cdb,dbf,afc->")
path, path_str = np.einsum_path(*edge_test3, optimize="greedy")
self.assert_path_equal(path, ["einsum_path", (1, 2), (0, 2), (0, 1)])
path, path_str = np.einsum_path(*edge_test3, optimize="optimal")
self.assert_path_equal(path, ["einsum_path", (1, 2), (0, 2), (0, 1)])
# Edge test4
edge_test4 = self.build_operands("dcc,fce,ea,dbf->ab")
path, path_str = np.einsum_path(*edge_test4, optimize="greedy")
self.assert_path_equal(path, ["einsum_path", (1, 2), (0, 1), (0, 1)])
path, path_str = np.einsum_path(*edge_test4, optimize="optimal")
self.assert_path_equal(path, ["einsum_path", (1, 2), (0, 2), (0, 1)])
# Edge test5
edge_test4 = self.build_operands("a,ac,ab,ad,cd,bd,bc->",
size_dict={
"a": 20,
"b": 20,
"c": 20,
"d": 20
})
path, path_str = np.einsum_path(*edge_test4, optimize="greedy")
self.assert_path_equal(path,
["einsum_path", (0, 1), (0, 1, 2, 3, 4, 5)])
path, path_str = np.einsum_path(*edge_test4, optimize="optimal")
self.assert_path_equal(path,
["einsum_path", (0, 1), (0, 1, 2, 3, 4, 5)])
def test_einsum_path(self):
# just check examples from np.einsum_path docstring
a = self.rng().rand(2, 2)
b = self.rng().rand(2, 5)
c = self.rng().rand(5, 2)
path_info = np.einsum_path('ij,jk,kl->il', a, b, c, optimize='greedy')
self.assertEqual(str(path_info[0]), "['einsum_path', (1, 2), (0, 1)]")
self.assertEqual(path_info[1].split('\n')[0],
' Complete contraction: ij,jk,kl->il')
# check this doesn't crash
I = self.rng().rand(10, 10, 10, 10)
C = self.rng().rand(10, 10)
np.einsum_path('ea,fb,abcd,gc,hd->efgh', C, C, I, C, C, optimize='greedy')
def path(self):
U1, U2, U1_, U2_ = [
unitary_group.rvs(4).reshape(2, 2, 2, 2) for _ in range(4)
]
M = unitary_group.rvs(2)
W = unitary_group.rvs(16).reshape(2, 2, 2, 2, 2, 2, 2, 2)
path = np.einsum_path(U2_, [6, 7, 26, 27],
U2_, [8, 9, 28, 29],
U2_, [10, 11, 30, 31],
U1_, [27, 28, 22, 23],
U1_, [29, 30, 24, 25],
W, [22, 23, 24, 25, 18, 19, 20, 21],
M, [26, 12],
M, [31, 17],
U1, [18, 19, 13, 14],
U1, [20, 21, 15, 16],
U2, [12, 13, 0, 1],
U2, [14, 15, 2, 3],
U2, [16, 17, 4, 5],
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
optimize="greedy")[0]
return path
def __call__(self, stage, itr):
"""!Record the single-particle correlators."""
S = self._inverter(stage, itr)
nx = S.shape[0]
nt = S.shape[1]
d = np.eye(nx*nt).reshape(*S.shape) # A kronecker delta
if self._roll is None:
self._roll = np.array([temporalRoller(nt, -t, fermionic=self.fermionic) for t in range(nt)])
# If there's no transformation needed, we should avoid doing
# space matrix-matrix-matrix, as it will scale poorly.
tensors = dict()
if self.transform is None:
tensors['destruction_creation'] = (self._roll, S)
tensors['creation_destruction'] = (self._roll, d-S)
else:
tensors['destruction_creation'] = (self._roll, self.transform.T.conj(), S, self.transform)
tensors['creation_destruction'] = (self._roll, self.transform.T.conj(), d-S, self.transform)
for c in self._einsum_paths:
if self._einsum_paths[c] is None:
self._einsum_paths[c], _ = np.einsum_path(self._indices[c], *tensors[c], optimize='optimal')
# The temporal roller sums over time, but does not *average* over time. So, divide by nt:
for name in self._einsum_paths:
res = self.nextItem(name)
np.einsum(self._indices[name], *tensors[name],
optimize=self._einsum_paths[name], out=res)
res /= nt
def kron_matvec(t, V):
'''matrix M multiply vectors V where M is kronecker product of A and B,
i.e. M = A\otimes B.
Args:
t (list of ndarray): M = tds[0] \otimes t[1] \otimes t[2] ...
V (ndarray): vectors to be multiplied.
Returns:
res (ndarray): results.
'''
shapes = [m.shape[1] for m in t]
if V.ndim == 1:
tmp = V.reshape(1, *shapes)
else:
tmp = V.reshape(V.shape[0], *shapes)
n = len(t)
params = []
for i, m in enumerate(t):
params.append(m)
params.append([i, n + i])
params.append(tmp)
params.append([2 * n] + list(range(n, 2 * n)))
params.append([2 * n] + list(range(n)))
path = np.einsum_path(*params, optimize='optimal')
res = np.einsum(*params, optimize=path[0])
res = res.reshape(res.shape[0], -1)
return res.squeeze()
def predict_slow(prj, x_tensors, path=None):
assert isinstance(x_tensors, list)
freq_dim = len(prj)
if freq_dim == 2:
subscripts = 'wrl,wrm, Ar,ABl,BCm,Co -> wo'
elif freq_dim == 3:
subscripts = 'wrl,wrm,wrn, Ar,ABl,BCm,CDn,Do -> wo'
else:
raise RuntimeError("Invalid freq_dim")
#subscripts = 'wrl,wrm, Dr, LMl, Mm'
assert len(x_tensors) == freq_dim+2
#print(type(prj), type(x_tensors))
operands = prj + x_tensors
if path is None:
from numpy import einsum_path
path, string_repr = einsum_path(subscripts, *operands, optimize=('greedy', 1E+18))
print(string_repr)
t1 = time.time()
t = numpy.einsum(subscripts, *operands, optimize=False)
t2 = time.time()
print(t2-t1, path)
return t, path
def _convert_color(imgs, same_size, method):
""" Helper function to convert colorspaces """
if method.endswith('gray'):
conversion = np.array([[0.0722], [0.7152], [0.2126]])
else:
conversion = np.array([[0.25, 0.5, 0.25], [-0.5, 0.0, 0.5],
[-0.25, 0.5, -0.25]])
if same_size:
path = 'greedy'
operation = 'bijk, kl -> bijl' if method.endswith(
'gray') else 'bijl, kl -> bijk'
else:
operation = 'ijk, kl -> ijl' if method.endswith(
'gray') else 'ijl, kl -> ijk'
path = np.einsum_path(operation,
imgs[0][..., :3],
conversion,
optimize='optimal')[0]
progress_bar = tqdm(imgs, desc="Converting", file=sys.stdout)
images = [
np.einsum(operation, img[..., :3], conversion,
optimize=path).astype('float32') for img in progress_bar
]
return images
def get_einsum(m, n):
# setup einsum_str
s0 = string.ascii_lowercase + string.ascii_uppercase
einsum_str = ''
for i in range(n):
einsum_str += s0[n + 2] + s0[i] + s0[i + 1] + ','
# remove last comma
einsum_str = einsum_str[:-1]
einsum_str += '->' + s0[n + 2] + s0[0] + s0[n]
# einsum path
test_array = np.zeros((n, m, 2, 2))
einsum_path = np.einsum_path(einsum_str, *test_array, optimize='greedy')
# calc matrix chain from m_n_2_2 tensor
# (If n > 2x alphabet length, einsum breaks -> split into k parts... n/k)
"""
1_n_(2x2)*...*1_3_(2x2)*1_2_(2x2)*1_1_(2x2) -> (2x2)_1
2_n_(2x2)*...*2_3_(2x2)*2_2_(2x2)*2_1_(2x2) -> (2x2)_2
.
.
m_n_(2x2)*...*m_3_(2x2)*m_2_(2x2)*m_1_(2x2) -> (2x2)_m
"""
def matrix_chain_calc(matrix_array):
return
return einsum_str, einsum_path
def EinsumGeneral(equation, *tensors, **kwargs):
tensors = list(tensors)
equation, isBinary = normalize_subscript(equation)
path = np.einsum_path(equation,
*[np.broadcast_to(np.nan, t.shape) for t in tensors],
**kwargs)
path = path[0][1:]
equation = equation.split('->')
eqs = equation[0].split(',')
target = equation[1]
for step in path:
if len(step) == 1:
result = EinsumFunction.apply(eqs[0] + '->' + target, tensors[0])
continue
assert step[0] < step[1]
in0 = tensors[step[0]]
in1 = tensors[step[1]]
tensors.pop(step[1])
tensors.pop(step[0])
tgt = _compute_target_tensor(eqs[step[0]], eqs[step[1]], target)
assert tgt != ""
eq = eqs[step[0]] + ',' + eqs[step[1]] + '->' + tgt
eqs.pop(step[1])
eqs.pop(step[0])
eqs.append(tgt)
result = EinsumFunction.apply(eq, in0, in1)
tensors.append(result)
return result
def einsum(*operands, **kwargs):
casting = kwargs.pop('casting', 'safe')
dtype = kwargs.pop('dtype', None)
optimize = kwargs.pop('optimize', False)
order = kwargs.pop('order', 'K')
split_every = kwargs.pop('split_every', None)
if kwargs:
raise TypeError("einsum() got unexpected keyword "
"argument(s) %s" % ",".join(kwargs))
einsum_dtype = dtype
inputs, outputs, ops = parse_einsum_input(operands)
subscripts = '->'.join((inputs, outputs))
# Infer the output dtype from operands
if dtype is None:
dtype = np.result_type(*[o.dtype for o in ops])
if einsum_can_optimize:
if optimize is not False:
# Avoid computation of dask arrays within np.einsum_path
# by passing in small numpy arrays broadcasted
# up to the right shape
fake_ops = [np.broadcast_to(o.dtype.type(0), shape=o.shape)
for o in ops]
optimize, _ = np.einsum_path(subscripts, *fake_ops,
optimize=optimize)
kwargs = {'optimize': optimize}
else:
kwargs = {}
inputs = [tuple(i) for i in inputs.split(",")]
# Set of all indices
all_inds = set(a for i in inputs for a in i)
# Which indices are contracted?
contract_inds = all_inds - set(outputs)
ncontract_inds = len(contract_inds)
# Introduce the contracted indices into the atop product
# so that we get numpy arrays, not lists
result = atop(chunk.einsum, tuple(outputs) + tuple(contract_inds),
*(a for ap in zip(ops, inputs) for a in ap),
# atop parameters
adjust_chunks={ind: 1 for ind in contract_inds}, dtype=dtype,
# np.einsum parameters
subscripts=subscripts, kernel_dtype=einsum_dtype,
ncontract_inds=ncontract_inds, order=order,
casting=casting, **kwargs)
# Now reduce over any extra contraction dimensions
if ncontract_inds > 0:
size = len(outputs)
return result.sum(axis=list(range(size, size + ncontract_inds)),
split_every=split_every)
return result
def einsum(*operands, dtype=None, optimize=False, split_every=None, **kwargs):
"""Dask added an additional keyword-only argument ``split_every``.
split_every: int >= 2 or dict(axis: int), optional
Determines the depth of the recursive aggregation.
Deafults to ``None`` which would let dask heuristically
decide a good default.
"""
einsum_dtype = dtype
inputs, outputs, ops = parse_einsum_input(operands)
subscripts = "->".join((inputs, outputs))
# Infer the output dtype from operands
if dtype is None:
dtype = np.result_type(*[o.dtype for o in ops])
if optimize is not False:
# Avoid computation of dask arrays within np.einsum_path
# by passing in small numpy arrays broadcasted
# up to the right shape
fake_ops = [np.broadcast_to(o.dtype.type(0), shape=o.shape) for o in ops]
optimize, _ = np.einsum_path(subscripts, *fake_ops, optimize=optimize)
inputs = [tuple(i) for i in inputs.split(",")]
# Set of all indices
all_inds = {a for i in inputs for a in i}
# Which indices are contracted?
contract_inds = all_inds - set(outputs)
ncontract_inds = len(contract_inds)
# Introduce the contracted indices into the blockwise product
# so that we get numpy arrays, not lists
result = blockwise(
chunk_einsum,
tuple(outputs) + tuple(contract_inds),
*(a for ap in zip(ops, inputs) for a in ap),
# blockwise parameters
adjust_chunks={ind: 1 for ind in contract_inds},
dtype=dtype,
# np.einsum parameters
subscripts=subscripts,
kernel_dtype=einsum_dtype,
ncontract_inds=ncontract_inds,
optimize=optimize,
**kwargs,
)
# Now reduce over any extra contraction dimensions
if ncontract_inds > 0:
size = len(outputs)
return result.sum(
axis=list(range(size, size + ncontract_inds)), split_every=split_every
)
return result
def test_memory_contraints(self):
# Ensure memory constraints are satisfied
outer_test = self.build_operands('a,b,c->abc')
path, path_str = np.einsum_path(*outer_test, optimize=('greedy', 0))
self.assert_path_equal(path, ['einsum_path', (0, 1, 2)])
path, path_str = np.einsum_path(*outer_test, optimize=('optimal', 0))
self.assert_path_equal(path, ['einsum_path', (0, 1, 2)])
long_test = self.build_operands('acdf,jbje,gihb,hfac')
path, path_str = np.einsum_path(*long_test, optimize=('greedy', 0))
self.assert_path_equal(path, ['einsum_path', (0, 1, 2, 3)])
path, path_str = np.einsum_path(*long_test, optimize=('optimal', 0))
self.assert_path_equal(path, ['einsum_path', (0, 1, 2, 3)])
def test_memory_contraints(self):
# Ensure memory constraints are satisfied
outer_test = self.build_operands('a,b,c->abc')
path, path_str = np.einsum_path(*outer_test, optimize=('greedy', 0))
self.assert_path_equal(path, ['einsum_path', (0, 1, 2)])
path, path_str = np.einsum_path(*outer_test, optimize=('optimal', 0))
self.assert_path_equal(path, ['einsum_path', (0, 1, 2)])
long_test = self.build_operands('acdf,jbje,gihb,hfac')
path, path_str = np.einsum_path(*long_test, optimize=('greedy', 0))
self.assert_path_equal(path, ['einsum_path', (0, 1, 2, 3)])
path, path_str = np.einsum_path(*long_test, optimize=('optimal', 0))
self.assert_path_equal(path, ['einsum_path', (0, 1, 2, 3)])
def backward(self, eta):
N, H, W, C = eta.shape
eta = eta.reshape(N, H//self.sacle, self.sacle, W//self.sacle, self.sacle, C)
if self.first_backward:
# 在第一次训练时计算优化路径
self.first_backward = False
self.backward_path = np.einsum_path('ijklmn->ijln', eta, optimize='greedy')[0]
return np.einsum('ijklmn->ijln', eta, optimize=True)
def make_decision_func(self, phi=None):
"""to compute the decision function, fully contract weights and input,
using W and Psi of the current instance.
Args:
phi (optional): ndarray, shape (number_of_inputs, N, 2),
alternative input tensor to use (e.g. for testing)
Returns:
ndarray, shape (number_of_inputs, label), decision functions for
given inputs
>>> f = self.make_decision_func()
>>> matches = np.argmax(f, axis=1)
"""
phi = phi if phi is not None else self.Phi
label_idx = self.N - 1
# fully contract components of MPS with input tensors
#
# shapes of involved quantities:
# W[i].shape = (2, bond_left, (label,) bond_right)
# Phi.shape = (N_train, N, 2)
# f_temp.shape = (t, bond_left, bond_right)
# f_final.shape = (t, label)
#
# note: `np.einsum` only allows 52 distinct indices
# starting point: contract first MPS component and its respective input
# tensor
f = np.einsum('dij,td->tij', self.W[0], phi[:, 0])
# step-wise contract each next MPS component and input tensor with the
# previous result
# exclude label tensor at first
path = np.einsum_path('tij,djk,td->tik',
f,
self.W[1],
phi[:, 1],
optimize='optimal')[0]
for i in range(1, label_idx):
f = np.einsum('tij,djk,td->tik',
f,
self.W[i],
phi[:, i],
optimize=path)
# contract with `label` tensor
# note the cyclic contraction over index `i`
f = np.einsum('tij,djli,td->tl',
f,
self.W[label_idx],
phi[:, label_idx],
optimize=True)
return f
def mo_ovlp_contraction(shape):
print("MO Overlaps")
_ = np.zeros(shape)
path, string_repr = np.einsum_path("pu,qv,uv->pq",
_, _, _,
optimize="optimal")
print(path)
print(string_repr)
print(SEP)
def backward(self, eta):
if self.is_test:
return eta
else:
if self.first_backward:
# 在第一次训练时计算优化路径
self.first_backward = False
self.backward_path = np.einsum_path('...,...->...', eta, self.mask, optimize='greedy')[0]
return np.einsum('...,...->...', eta, self.mask, optimize=self.backward_path)
def test_path_type_input(self):
# Test explicit path handeling
path_test = self.build_operands('dcc,fce,ea,dbf->ab')
path, path_str = np.einsum_path(*path_test, optimize=False)
self.assert_path_equal(path, ['einsum_path', (0, 1, 2, 3)])
path, path_str = np.einsum_path(*path_test, optimize=True)
self.assert_path_equal(path, ['einsum_path', (1, 2), (0, 1), (0, 1)])
exp_path = ['einsum_path', (0, 2), (0, 2), (0, 1)]
path, path_str = np.einsum_path(*path_test, optimize=exp_path)
self.assert_path_equal(path, exp_path)
# Double check einsum works on the input path
noopt = np.einsum(*path_test, optimize=False)
opt = np.einsum(*path_test, optimize=exp_path)
assert_almost_equal(noopt, opt)
def test_path_type_input(self):
# Test explicit path handeling
path_test = self.build_operands('dcc,fce,ea,dbf->ab')
path, path_str = np.einsum_path(*path_test, optimize=False)
self.assert_path_equal(path, ['einsum_path', (0, 1, 2, 3)])
path, path_str = np.einsum_path(*path_test, optimize=True)
self.assert_path_equal(path, ['einsum_path', (0, 3), (0, 2), (0, 1)])
exp_path = ['einsum_path', (0, 2), (0, 2), (0, 1)]
path, path_str = np.einsum_path(*path_test, optimize=exp_path)
self.assert_path_equal(path, exp_path)
# Double check einsum works on the input path
noopt = np.einsum(*path_test, optimize=False)
opt = np.einsum(*path_test, optimize=exp_path)
assert_almost_equal(noopt, opt)
def test_optimal_edge_cases():
# Edge test5
expression = 'a,ac,ab,ad,cd,bd,bc->'
edge_test4 = oe.helpers.build_views(expression,
dimension_dict={
"a": 20,
"b": 20,
"c": 20,
"d": 20
})
path, path_str = np.einsum_path(expression, *edge_test4, optimize='greedy')
check_path(path, ['einsum_path', (0, 1), (0, 1, 2, 3, 4, 5)])
path, path_str = np.einsum_path(expression,
*edge_test4,
optimize='optimal')
check_path(path, ['einsum_path', (0, 1), (0, 1, 2, 3, 4, 5)])
def einsum(*operands, **kwargs):
dtype = kwargs.pop("dtype", None)
optimize = kwargs.pop("optimize", False)
split_every = kwargs.pop("split_every", None)
einsum_dtype = dtype
inputs, outputs, ops = parse_einsum_input(operands)
subscripts = "->".join((inputs, outputs))
# Infer the output dtype from operands
if dtype is None:
dtype = np.result_type(*[o.dtype for o in ops])
if optimize is not False:
# Avoid computation of dask arrays within np.einsum_path
# by passing in small numpy arrays broadcasted
# up to the right shape
fake_ops = [
np.broadcast_to(o.dtype.type(0), shape=o.shape) for o in ops
]
optimize, _ = np.einsum_path(subscripts, *fake_ops, optimize=optimize)
inputs = [tuple(i) for i in inputs.split(",")]
# Set of all indices
all_inds = {a for i in inputs for a in i}
# Which indices are contracted?
contract_inds = all_inds - set(outputs)
ncontract_inds = len(contract_inds)
# Introduce the contracted indices into the blockwise product
# so that we get numpy arrays, not lists
result = blockwise(
chunk_einsum,
tuple(outputs) + tuple(contract_inds),
*(a for ap in zip(ops, inputs) for a in ap),
# blockwise parameters
adjust_chunks={ind: 1
for ind in contract_inds},
dtype=dtype,
# np.einsum parameters
subscripts=subscripts,
kernel_dtype=einsum_dtype,
ncontract_inds=ncontract_inds,
optimize=optimize,
**kwargs,
)
# Now reduce over any extra contraction dimensions
if ncontract_inds > 0:
size = len(outputs)
return result.sum(axis=list(range(size, size + ncontract_inds)),
split_every=split_every)
return result
def einsum(*operands, **kwargs):
dtype = kwargs.get('dtype')
optimize = kwargs.get('optimize')
einsum_dtype = dtype
inputs, outputs, ops = parse_einsum_input(operands)
subscripts = '->'.join((inputs, outputs))
# Infer the output dtype from operands
if dtype is None:
dtype = np.result_type(*[o.dtype for o in ops])
if optimize is None:
optimize = False
if einsum_can_optimize and optimize is not False:
# Avoid computation of dask arrays within np.einsum_path
# by passing in small numpy arrays broadcasted
# up to the right shape
fake_ops = [np.broadcast_to(o.dtype.type(0), shape=o.shape)
for o in ops]
optimize, _ = np.einsum_path(subscripts, *fake_ops, optimize=optimize)
inputs = [tuple(i) for i in inputs.split(",")]
# Set of all indices
all_inds = set(a for i in inputs for a in i)
# Which indices are contracted?
contract_inds = all_inds - set(outputs)
ncontract_inds = len(contract_inds)
# Update kwargs with np.einsum parameters
kwargs['subscripts'] = subscripts
kwargs['kernel_dtype'] = einsum_dtype
kwargs['ncontract_inds'] = ncontract_inds
if einsum_can_optimize:
kwargs['optimize'] = optimize
# Update kwargs with atop parameters
kwargs['adjust_chunks'] = {ind: 1 for ind in contract_inds}
kwargs['dtype'] = dtype
# Introduce the contracted indices into the atop product
# so that we get numpy arrays, not lists
result = atop(_einsum_kernel, tuple(outputs) + tuple(contract_inds),
*(a for ap in zip(ops, inputs) for a in ap),
**kwargs)
# Now reduce over any extra contraction dimensions
if ncontract_inds > 0:
size = len(outputs)
return result.sum(axis=list(range(size, size + ncontract_inds)))
return result
def test_muti_operand(self):
e = np.random.rand(50, 10)
f = np.random.rand(50, 10, 10)
g = np.random.rand(50, 10, 20)
h = np.random.rand(20, 20, 10)
subs = '...k,...km,...kp,plo->...lom'
path = np.einsum_path(subs, e, f, g, h)
res0 = np.einsum(subs, e, f, g, h, optimize=path[0])
res1 = einsumt(subs, e, f, g, h, optimize=path[0])
assert np.allclose(res0, res1)
def _generate_efms(self):
self.efm_einstrs, self.efm_specs, self.efm_einpaths = [], [], []
if self.gen_efms:
for edgs,ws in zip(self.edges, self.weights):
einstr, efm_spec = efp2efms(EFP(edgs, weights=ws).graph)
self.efm_einstrs.append(einstr)
self.efm_specs.append(efm_spec)
self.efm_einpaths.append(np.einsum_path(einstr,
*[np.empty([4]*sum(s)) for s in efm_spec],
optimize=self.ve.np_optimize)[0])
def forward(self, x):
if self.is_test:
# 如果是测试，直接返回
return x
else:
# 如果是训练，按概率将输出置为0
self.mask = np.random.uniform(0, 1, x.shape) > self.drop_rate
if self.first_forward:
# 在第一次训练时计算优化路径
self.first_forward = False
self.forward_path = np.einsum_path('...,...,->...', x, self.mask, self.fix_value, optimize='greedy')[0]
return np.einsum('...,...,->...', x, self.mask, self.fix_value, optimize=self.first_forward)
def coeff_ovlp_contraction(b, k):
bra_coeffs = np.zeros(b)
ket_coeffs = np.zeros(k)
alpha_ovlps = np.zeros((b, k))
beta_ovlps = np.zeros((b, k))
path, string_repr = np.einsum_path("b,k,bk,bk",
bra_coeffs, ket_coeffs, alpha_ovlps, beta_ovlps,
optimize="optimal")
print(path)
print(string_repr)
print(SEP)
def test_einsum_path(self):
# Use the examples from the numpy docs.
npa = np.random.rand(2, 2)
npb = np.random.rand(2, 5)
npc = np.random.rand(5, 2)
a = array.from_buffer(npa)
b = array.from_buffer(npb)
c = array.from_buffer(npc)
path_info = np.einsum_path('ij,jk,kl->il', a, b, c, optimize='greedy')
self.assertEqual(path_info[0], ['einsum_path', (1, 2), (0, 1)])
def _getEinsumPath(self, path, d):
"""!
Return an optimized einsum path.
"""
try:
return self._einsum_paths[path]
except KeyError:
self._einsum_paths[path], _ = np.einsum_path(path,
d,
d,
optimize="optimal")
getLogger(__name__).info("Optimized Einsum path %s", path)
return self._einsum_paths[path]
def test_long_paths(self):
# Long complex cases
# Long test 1
long_test1 = self.build_operands('acdf,jbje,gihb,hfac,gfac,gifabc,hfac')
path, path_str = np.einsum_path(*long_test1, optimize='greedy')
self.assert_path_equal(path, ['einsum_path',
(1, 4), (2, 4), (1, 4), (1, 3), (1, 2), (0, 1)])
path, path_str = np.einsum_path(*long_test1, optimize='optimal')
self.assert_path_equal(path, ['einsum_path',
(3, 6), (3, 4), (2, 4), (2, 3), (0, 2), (0, 1)])
# Long test 2
long_test2 = self.build_operands('chd,bde,agbc,hiad,bdi,cgh,agdb')
path, path_str = np.einsum_path(*long_test2, optimize='greedy')
self.assert_path_equal(path, ['einsum_path',
(3, 4), (0, 3), (3, 4), (1, 3), (1, 2), (0, 1)])
path, path_str = np.einsum_path(*long_test2, optimize='optimal')
self.assert_path_equal(path, ['einsum_path',
(0, 5), (1, 4), (3, 4), (1, 3), (1, 2), (0, 1)])
def __init__(self, psi_0, op, e_vals, e_states):
self.C = psi_0 @ e_states
energy_differences = np.tile(e_vals, (e_vals.shape[0], 1))
self.energy_diff_array = energy_differences - energy_differences.T
self.eev_including_off_diagonal = e_states.T @ op @ e_states
self.phases = np.zeros_like(self.energy_diff_array,
dtype=np.complex128)
self.einsum_path = np.einsum_path('a, b, ab, ab -> ',
self.C,
self.C,
self.phases,
self.eev_including_off_diagonal,
optimize='optimal')[0]