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 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]