def test_sharing_modulo_commutativity(eq, backend): ops = helpers.build_views(eq) ops = [to_backend[backend](x) for x in ops] inputs, output, _ = parse_einsum_input([eq] + ops) inputs = inputs.split(',') print('-' * 40) print('Without sharing:') with shared_intermediates() as cache: _einsum(eq, *ops, backend=backend) expected = count_cached_ops(cache) print('-' * 40) print('With sharing:') with shared_intermediates() as cache: for permuted in itertools.permutations(zip(inputs, ops)): permuted_inputs = [p[0] for p in permuted] permuted_ops = [p[1] for p in permuted] permuted_eq = '{}->{}'.format(','.join(permuted_inputs), output) _einsum(permuted_eq, *permuted_ops, backend=backend) actual = count_cached_ops(cache) print('-' * 40) print('Without sharing: {} expressions'.format(expected)) print('With sharing: {} expressions'.format(actual)) assert actual == expected
def test_sparse(string): views = helpers.build_views(string) # sparsify views so they don't become dense during contraction for view in views: np.random.seed(42) mask = np.random.choice([False, True], view.shape, True, [0.05, 0.95]) view[mask] = 0 ein = contract(string, *views, optimize=False, use_blas=False) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) # test non-conversion mode sparse_views = [sparse.COO.from_numpy(x) for x in views] sparse_opt = expr(*sparse_views, backend='sparse') # check type is maintained when not using numpy arrays assert isinstance(sparse_opt, sparse.COO) assert np.allclose(ein, sparse_opt.todense()) # try raw contract sparse_opt = contract(string, *sparse_views, backend='sparse') assert isinstance(sparse_opt, sparse.COO) assert np.allclose(ein, sparse_opt.todense())
def test_partial_sharing(backend): eq = 'ab,bc,de->' x, y, z1 = helpers.build_views(eq) z2 = 2.0 * z1 - 1.0 expr = contract_expression(eq, x.shape, y.shape, z1.shape) print('-' * 40) print('Without sharing:') num_exprs_nosharing = Counter() with shared_intermediates() as cache: expr(x, y, z1, backend=backend) num_exprs_nosharing.update(count_cached_ops(cache)) with shared_intermediates() as cache: expr(x, y, z2, backend=backend) num_exprs_nosharing.update(count_cached_ops(cache)) print('-' * 40) print('With sharing:') with shared_intermediates() as cache: expr(x, y, z1, backend=backend) expr(x, y, z2, backend=backend) num_exprs_sharing = count_cached_ops(cache) print('-' * 40) print('Without sharing: {} expressions'.format(num_exprs_nosharing)) print('With sharing: {} expressions'.format(num_exprs_sharing)) assert num_exprs_nosharing['einsum'] > num_exprs_sharing['einsum']
def test_sharing_modulo_commutativity(eq, backend): ops = helpers.build_views(eq) ops = [to_backend[backend](x) for x in ops] inputs, output, _ = parse_einsum_input([eq] + ops) inputs = inputs.split(',') print('-' * 40) print('Without sharing:') with shared_intermediates() as cache: _einsum(eq, *ops, backend=backend) expected = count_cached_ops(cache) print('-' * 40) print('With sharing:') with shared_intermediates() as cache: for permuted in itertools.permutations(zip(inputs, ops)): permuted_inputs = [p[0] for p in permuted] permuted_ops = [p[1] for p in permuted] permuted_eq = '{}->{}'.format(','.join(permuted_inputs), output) _einsum(permuted_eq, *permuted_ops, backend=backend) actual = count_cached_ops(cache) print('-' * 40) print('Without sharing: {} expressions'.format(expected)) print('With sharing: {} expressions'.format(actual)) assert actual == expected
def test_partial_sharing(backend): eq = 'ab,bc,de->' x, y, z1 = helpers.build_views(eq) z2 = 2.0 * z1 - 1.0 expr = contract_expression(eq, x.shape, y.shape, z1.shape) print('-' * 40) print('Without sharing:') num_exprs_nosharing = Counter() with shared_intermediates() as cache: expr(x, y, z1, backend=backend) num_exprs_nosharing.update(count_cached_ops(cache)) with shared_intermediates() as cache: expr(x, y, z2, backend=backend) num_exprs_nosharing.update(count_cached_ops(cache)) print('-' * 40) print('With sharing:') with shared_intermediates() as cache: expr(x, y, z1, backend=backend) expr(x, y, z2, backend=backend) num_exprs_sharing = count_cached_ops(cache) print('-' * 40) print('Without sharing: {} expressions'.format(num_exprs_nosharing)) print('With sharing: {} expressions'.format(num_exprs_sharing)) assert num_exprs_nosharing['einsum'] > num_exprs_sharing['einsum']
def test_no_sharing_separate_cache(backend): eq = 'ab,bc,cd->' views = helpers.build_views(eq) expr = contract_expression(eq, *(v.shape for v in views)) print('-' * 40) print('Without sharing:') with shared_intermediates() as cache: expr(*views, backend=backend) expected = count_cached_ops(cache) expected.update(count_cached_ops(cache)) # we expect double print('-' * 40) print('With sharing:') with shared_intermediates() as cache1: expr(*views, backend=backend) actual = count_cached_ops(cache1) with shared_intermediates() as cache2: expr(*views, backend=backend) actual.update(count_cached_ops(cache2)) print('-' * 40) print('Without sharing: {} expressions'.format(expected)) print('With sharing: {} expressions'.format(actual)) assert actual == expected
def test_no_sharing_separate_cache(backend): eq = 'ab,bc,cd->' views = helpers.build_views(eq) expr = contract_expression(eq, *(v.shape for v in views)) print('-' * 40) print('Without sharing:') with shared_intermediates() as cache: expr(*views, backend=backend) expected = count_cached_ops(cache) expected.update(count_cached_ops(cache)) # we expect double print('-' * 40) print('With sharing:') with shared_intermediates() as cache1: expr(*views, backend=backend) actual = count_cached_ops(cache1) with shared_intermediates() as cache2: expr(*views, backend=backend) actual.update(count_cached_ops(cache2)) print('-' * 40) print('Without sharing: {} expressions'.format(expected)) print('With sharing: {} expressions'.format(actual)) assert actual == expected
def test_rand_equation(optimize, n, reg, n_out, global_dim): eq, _, size_dict = helpers.rand_equation(n, reg, n_out, d_min=2, d_max=5, seed=42, return_size_dict=True) views = helpers.build_views(eq, size_dict) expected = contract(eq, *views, optimize=False) actual = contract(eq, *views, optimize=optimize) assert np.allclose(expected, actual)
def test_rand_equation(optimize, n, reg, n_out, global_dim): eq, _, size_dict = helpers.rand_equation(n, reg, n_out, d_min=2, d_max=5, seed=42, return_size_dict=True) views = helpers.build_views(eq, size_dict) expected = contract(eq, *views, optimize=False) actual = contract(eq, *views, optimize=optimize) assert np.allclose(expected, actual)
def fn(): X, Y, Z = helpers.build_views('ab,bc,cd') with shared_intermediates(): contract('ab,bc,cd->a', X, Y, Z) contract('ab,bc,cd->b', X, Y, Z) return len(get_sharing_cache())
def fn(): X, Y, Z = helpers.build_views('ab,bc,cd') with shared_intermediates(): contract('ab,bc,cd->a', X, Y, Z) contract('ab,bc,cd->b', X, Y, Z) return len(get_sharing_cache())
def test_contract_expressions(string, optimize, use_blas, out_spec): views = helpers.build_views(string) shapes = [view.shape for view in views] expected = contract(string, *views, optimize=False, use_blas=False) expr = contract_expression(string, *shapes, optimize=optimize, use_blas=use_blas) if out_spec and ("->" in string) and (string[-2:] != "->"): out, = helpers.build_views(string.split('->')[1]) expr(*views, out=out) else: out = expr(*views) assert np.allclose(out, expected) # check representations assert string in expr.__repr__() assert string in expr.__str__()
def test_contract_expressions(string, optimize, use_blas, out_spec): views = helpers.build_views(string) shapes = [view.shape for view in views] expected = contract(string, *views, optimize=False, use_blas=False) expr = contract_expression(string, *shapes, optimize=optimize, use_blas=use_blas) if out_spec and ("->" in string) and (string[-2:] != "->"): out, = helpers.build_views(string.split('->')[1]) expr(*views, out=out) else: out = expr(*views) assert np.allclose(out, expected) # check representations assert string in expr.__repr__() assert string in expr.__str__()
def test_compare_blas(string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False) opt = contract(string, *views, optimize='greedy') assert np.allclose(ein, opt) opt = contract(string, *views, optimize='optimal') assert np.allclose(ein, opt)
def test_sharing_value(eq, backend): views = helpers.build_views(eq) shapes = [v.shape for v in views] expr = contract_expression(eq, *shapes) expected = expr(*views, backend=backend) with shared_intermediates(): actual = expr(*views, backend=backend) assert (actual == expected).all()
def test_jax(string): # pragma: no cover views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) opt = expr(*views, backend='jax') assert np.allclose(ein, opt) assert isinstance(opt, np.ndarray)
def test_sharing_value(eq, backend): views = helpers.build_views(eq) shapes = [v.shape for v in views] expr = contract_expression(eq, *shapes) expected = expr(*views, backend=backend) with shared_intermediates(): actual = expr(*views, backend=backend) assert (actual == expected).all()
def test_compare_blas_greek(optimize, string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False) # convert to greek string = ''.join(chr(ord(c) + 848) if c not in ',->.' else c for c in string) opt = contract(string, *views, optimize=optimize) assert np.allclose(ein, opt)
def test_compare_blas_greek(optimize, string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False) # convert to greek string = ''.join(compat.get_char(ord(c) + 848) if c not in ',->.' else c for c in string) opt = contract(string, *views, optimize=optimize) assert np.allclose(ein, opt)
def test_theano(string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) opt = expr(*views, backend='theano') assert np.allclose(ein, opt) # test non-conversion mode theano_views = backends.convert_arrays_to_theano(views) theano_opt = expr(*theano_views, backend='theano') assert isinstance(theano_opt, theano.tensor.TensorVariable)
def test_cupy(string): # pragma: no cover views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) opt = expr(*views, backend='cupy') assert np.allclose(ein, opt) # test non-conversion mode cupy_views = backends.convert_arrays_to_cupy(views) cupy_opt = expr(*cupy_views, backend='cupy') assert isinstance(cupy_opt, cupy.ndarray) assert np.allclose(ein, cupy.asnumpy(cupy_opt))
def test_torch(string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) opt = expr(*views, backend='torch') assert np.allclose(ein, opt) # test non-conversion mode torch_views = [backends.to_torch(view) for view in views] torch_opt = expr(*torch_views) assert isinstance(torch_opt, torch.Tensor) assert np.allclose(ein, torch_opt.cpu().numpy())
def test_tensorflow(string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) opt = np.empty_like(ein) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) sess = tf.Session() with sess.as_default(): expr(*views, backend='tensorflow', out=opt) assert np.allclose(ein, opt) # test non-conversion mode tensorflow_views = backends.convert_arrays_to_tensorflow(views) expr(*tensorflow_views, backend='tensorflow')
def test_contract_expression_with_constants(string, constants): views = helpers.build_views(string) expected = contract(string, *views, optimize=False, use_blas=False) shapes = [view.shape for view in views] expr_args = [] ctrc_args = [] for i, (shape, view) in enumerate(zip(shapes, views)): if i in constants: expr_args.append(view) else: expr_args.append(shape) ctrc_args.append(view) expr = contract_expression(string, *expr_args, constants=constants) print(expr) out = expr(*ctrc_args) assert np.allclose(expected, out)
def test_dask(string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) # test non-conversion mode da_views = [da.from_array(x, chunks=(2)) for x in views] da_opt = expr(*da_views, backend='dask') # check type is maintained when not using numpy arrays assert isinstance(da_opt, da.Array) assert np.allclose(ein, np.array(da_opt)) # try raw contract da_opt = contract(string, *da_views, backend='dask') assert isinstance(da_opt, da.Array) assert np.allclose(ein, np.array(da_opt))
def test_contract_expression_with_constants(string, constants): views = helpers.build_views(string) expected = contract(string, *views, optimize=False, use_blas=False) shapes = [view.shape for view in views] expr_args = [] ctrc_args = [] for i, (shape, view) in enumerate(zip(shapes, views)): if i in constants: expr_args.append(view) else: expr_args.append(shape) ctrc_args.append(view) expr = contract_expression(string, *expr_args, constants=constants) print(expr) out = expr(*ctrc_args) assert np.allclose(expected, out)
def test_tensor_blas(inp, benchmark): # Weed out non-blas cases if benchmark is False: return tensor_strs, output, reduced_idx = inp einsum_str = ','.join(tensor_strs) + '->' + output # Only binary operations should be here if len(tensor_strs) != 2: assert False view_left, view_right = helpers.build_views(einsum_str) einsum_result = np.einsum(einsum_str, view_left, view_right) blas_result = blas.tensor_blas(view_left, tensor_strs[0], view_right, tensor_strs[1], output, reduced_idx) assert np.allclose(einsum_result, blas_result)
def test_object_arrays_backend(string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) assert ein.dtype != object shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) obj_views = [view.astype(object) for view in views] # try raw contract obj_opt = contract(string, *obj_views, backend="object") assert obj_opt.dtype == object assert np.allclose(ein, obj_opt.astype(float)) # test expression obj_opt = expr(*obj_views, backend="object") assert obj_opt.dtype == object assert np.allclose(ein, obj_opt.astype(float))
def test_theano_with_sharing(string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) with sharing.shared_intermediates() as cache: thn1 = expr(*views, backend='theano') assert sharing.get_sharing_cache() is cache cache_sz = len(cache) assert cache_sz > 0 thn2 = expr(*views, backend='theano') assert len(cache) == cache_sz assert all( isinstance(t, theano.tensor.TensorVariable) for t in cache.values()) assert np.allclose(ein, thn1) assert np.allclose(ein, thn2)
def test_tensor_blas(inp, benchmark): # Weed out non-blas cases if benchmark is False: return tensor_strs, output, reduced_idx = inp einsum_str = ','.join(tensor_strs) + '->' + output # Only binary operations should be here if len(tensor_strs) != 2: assert False view_left, view_right = helpers.build_views(einsum_str) einsum_result = np.einsum(einsum_str, view_left, view_right) blas_result = blas.tensor_blas(view_left, tensor_strs[0], view_right, tensor_strs[1], output, reduced_idx) assert np.allclose(einsum_result, blas_result)
def test_tensorflow_with_sharing(string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) sess = tf.Session(config=_TF_CONFIG) with sess.as_default(), sharing.shared_intermediates() as cache: tfl1 = expr(*views, backend='tensorflow') assert sharing.get_sharing_cache() is cache cache_sz = len(cache) assert cache_sz > 0 tfl2 = expr(*views, backend='tensorflow') assert len(cache) == cache_sz assert all(isinstance(t, tf.Tensor) for t in cache.values()) assert np.allclose(ein, tfl1) assert np.allclose(ein, tfl2)
def test_complete_sharing(backend): eq = 'ab,bc,cd->' views = helpers.build_views(eq) expr = contract_expression(eq, *(v.shape for v in views)) print('-' * 40) print('Without sharing:') with shared_intermediates() as cache: expr(*views, backend=backend) expected = count_cached_ops(cache) print('-' * 40) print('With sharing:') with shared_intermediates() as cache: expr(*views, backend=backend) expr(*views, backend=backend) actual = count_cached_ops(cache) print('-' * 40) print('Without sharing: {} expressions'.format(expected)) print('With sharing: {} expressions'.format(actual)) assert actual == expected
def test_complete_sharing(backend): eq = 'ab,bc,cd->' views = helpers.build_views(eq) expr = contract_expression(eq, *(v.shape for v in views)) print('-' * 40) print('Without sharing:') with shared_intermediates() as cache: expr(*views, backend=backend) expected = count_cached_ops(cache) print('-' * 40) print('With sharing:') with shared_intermediates() as cache: expr(*views, backend=backend) expr(*views, backend=backend) actual = count_cached_ops(cache) print('-' * 40) print('Without sharing: {} expressions'.format(expected)) print('With sharing: {} expressions'.format(actual)) assert actual == expected
def test_sharing_nesting(backend): eqs = ['ab,bc,cd->a', 'ab,bc,cd->b', 'ab,bc,cd->c', 'ab,bc,cd->c'] views = helpers.build_views(eqs[0]) shapes = [v.shape for v in views] refs = weakref.WeakValueDictionary() def method1(views): with shared_intermediates(): w = contract_expression(eqs[0], *shapes)(*views, backend=backend) x = contract_expression(eqs[2], *shapes)(*views, backend=backend) result = contract_expression('a,b->', w.shape, x.shape)(w, x, backend=backend) refs['w'] = w refs['x'] = x del w, x assert 'w' in refs assert 'x' in refs assert 'w' not in refs, 'cache leakage' assert 'x' not in refs, 'cache leakage' return result def method2(views): with shared_intermediates(): y = contract_expression(eqs[2], *shapes)(*views, backend=backend) z = contract_expression(eqs[3], *shapes)(*views, backend=backend) refs['y'] = y refs['z'] = z result = contract_expression('c,d->', y.shape, z.shape)(y, z, backend=backend) result = result + method1(views) # nest method1 in method2 del y, z assert 'y' in refs assert 'z' in refs assert 'y' not in refs assert 'z' not in refs method1(views) method2(views)
def test_sharing_nesting(backend): eqs = ["ab,bc,cd->a", "ab,bc,cd->b", "ab,bc,cd->c", "ab,bc,cd->c"] views = helpers.build_views(eqs[0]) shapes = [v.shape for v in views] refs = weakref.WeakValueDictionary() def method1(views): with shared_intermediates(): w = contract_expression(eqs[0], *shapes)(*views, backend=backend) x = contract_expression(eqs[2], *shapes)(*views, backend=backend) result = contract_expression("a,b->", w.shape, x.shape)(w, x, backend=backend) refs["w"] = w refs["x"] = x del w, x assert "w" in refs assert "x" in refs assert "w" not in refs, "cache leakage" assert "x" not in refs, "cache leakage" return result def method2(views): with shared_intermediates(): y = contract_expression(eqs[2], *shapes)(*views, backend=backend) z = contract_expression(eqs[3], *shapes)(*views, backend=backend) refs["y"] = y refs["z"] = z result = contract_expression("c,d->", y.shape, z.shape)(y, z, backend=backend) result = result + method1(views) # nest method1 in method2 del y, z assert "y" in refs assert "z" in refs assert "y" not in refs assert "z" not in refs method1(views) method2(views)
def test_sharing_nesting(backend): eqs = ['ab,bc,cd->a', 'ab,bc,cd->b', 'ab,bc,cd->c', 'ab,bc,cd->c'] views = helpers.build_views(eqs[0]) shapes = [v.shape for v in views] refs = weakref.WeakValueDictionary() def method1(views): with shared_intermediates(): w = contract_expression(eqs[0], *shapes)(*views, backend=backend) x = contract_expression(eqs[2], *shapes)(*views, backend=backend) result = contract_expression('a,b->', w.shape, x.shape)(w, x, backend=backend) refs['w'] = w refs['x'] = x del w, x assert 'w' in refs assert 'x' in refs assert 'w' not in refs, 'cache leakage' assert 'x' not in refs, 'cache leakage' return result def method2(views): with shared_intermediates(): y = contract_expression(eqs[2], *shapes)(*views, backend=backend) z = contract_expression(eqs[3], *shapes)(*views, backend=backend) refs['y'] = y refs['z'] = z result = contract_expression('c,d->', y.shape, z.shape)(y, z, backend=backend) result = result + method1(views) # nest method1 in method2 del y, z assert 'y' in refs assert 'z' in refs assert 'y' not in refs assert 'z' not in refs method1(views) method2(views)
def test_sharing_reused_cache(backend): eq = "ab,bc,cd->" views = helpers.build_views(eq) expr = contract_expression(eq, *(v.shape for v in views)) print("-" * 40) print("Without sharing:") with shared_intermediates() as cache: expr(*views, backend=backend) expected = count_cached_ops(cache) print("-" * 40) print("With sharing:") with shared_intermediates() as cache: expr(*views, backend=backend) with shared_intermediates(cache): expr(*views, backend=backend) actual = count_cached_ops(cache) print("-" * 40) print("Without sharing: {} expressions".format(expected)) print("With sharing: {} expressions".format(actual)) assert actual == expected
def test_linear_vs_ssa(equation): views = helpers.build_views(equation) linear_path, _ = contract_path(equation, *views) ssa_path = linear_to_ssa(linear_path) linear_path2 = ssa_to_linear(ssa_path) assert linear_path2 == linear_path
def test_printing(): string = "bbd,bda,fc,db->acf" views = helpers.build_views(string) ein = contract_path(string, *views) assert len(str(ein[1])) == 728
def test_drop_in_replacement(string): views = helpers.build_views(string) opt = contract(string, *views) assert np.allclose(opt, np.einsum(string, *views))
def test_compare(optimize, string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) opt = contract(string, *views, optimize=optimize, use_blas=False) assert np.allclose(ein, opt)
def test_drop_in_replacement(string): views = helpers.build_views(string) opt = contract(string, *views) assert np.allclose(opt, np.einsum(string, *views))
def test_compare_blas(optimize, string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False) opt = contract(string, *views, optimize=optimize) assert np.allclose(ein, opt)
def test_linear_vs_ssa(equation): views = helpers.build_views(equation) linear_path, _ = contract_path(equation, *views) ssa_path = linear_to_ssa(linear_path) linear_path2 = ssa_to_linear(ssa_path) assert linear_path2 == linear_path
def test_printing(): string = "bbd,bda,fc,db->acf" views = helpers.build_views(string) ein = contract_path(string, *views) assert len(str(ein[1])) == 726