Trajectree 0.0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

trajectree/quimb/tests/test_tensor/test_contract.py

@@ -0,0 +1,67 @@
+import pytest
+import numpy as np
+
+import quimb.tensor as qtn
+from quimb.tensor.contraction import _CONTRACT_BACKEND, _TENSOR_LINOP_BACKEND
+
+
+class TestContractOpts:
+    def test_contract_strategy(self):
+        assert qtn.get_contract_strategy() == "greedy"
+        with qtn.contract_strategy("auto"):
+            assert qtn.get_contract_strategy() == "auto"
+        assert qtn.get_contract_strategy() == "greedy"
+
+    def test_contract_backend(self):
+        assert qtn.get_contract_backend() == _CONTRACT_BACKEND
+        with qtn.contract_backend("cupy"):
+            assert qtn.get_contract_backend() == "cupy"
+        assert qtn.get_contract_backend() == _CONTRACT_BACKEND
+
+    def test_tensor_linop_backend(self):
+        assert qtn.get_tensor_linop_backend() == _TENSOR_LINOP_BACKEND
+        with qtn.tensor_linop_backend("cupy"):
+            assert qtn.get_tensor_linop_backend() == "cupy"
+        assert qtn.get_tensor_linop_backend() == _TENSOR_LINOP_BACKEND
+
+    def test_contract_cache(self):
+        import cotengra as ctg
+
+        info = {"num_calls": 0}
+
+        def my_custom_opt(inputs, output, size_dict, memory_limit=None):
+            info["num_calls"] += 1
+            return [(0, 1)] * (len(inputs) - 1)
+
+        ctg.register_preset("quimb_test_opt", my_custom_opt)
+
+        tn = qtn.MPS_rand_state(4, 3) & qtn.MPS_rand_state(4, 3)
+        assert tn.contract(
+            all, optimize="quimb_test_opt", get="expression"
+        ) is tn.contract(all, optimize="quimb_test_opt", get="expression")
+        assert info["num_calls"] == 1
+
+        assert info["num_calls"] == 1
+
+
+@pytest.mark.parametrize("around", ["I3,3", "I0,0", "I1,2"])
+@pytest.mark.parametrize("equalize_norms", [False, True])
+@pytest.mark.parametrize("gauge_boundary_only", [False, True])
+def test_contract_approx_with_gauges(
+    around, equalize_norms, gauge_boundary_only
+):
+    rng = np.random.default_rng(42)
+    tn = qtn.TN2D_from_fill_fn(
+        lambda shape: rng.uniform(size=shape, low=-0.5), 7, 7, 4
+    )
+    Zex = tn ^ ...
+    Z = tn.contract_around(
+        around,
+        max_bond=8,
+        gauges=True,
+        gauge_boundary_only=gauge_boundary_only,
+        tree_gauge_distance=2,
+        equalize_norms=equalize_norms,
+        progbar=True,
+    )
+    assert Z == pytest.approx(Zex, rel=1e-2)

trajectree/quimb/tests/test_tensor/test_decomp.py

@@ -0,0 +1,40 @@
+import pytest
+import numpy as np
+import quimb as qu
+
+
+def test_sgn_convention():
+    from quimb.tensor.decomp import sgn
+
+    assert sgn(1) == 1
+    assert sgn(2.0) == 1
+    assert sgn(-1) == -1
+    assert sgn(-2.0) == -1
+    assert sgn(0) == 1
+    assert sgn(0.0) == 1
+    assert sgn(0.0 + 0.0j) == 1
+    assert sgn(1.0 + 2.0j) != 1
+    assert sgn(1.0 + 2.0j) != -1
+    assert abs(sgn(1.0 + 2.0j)) == pytest.approx(1)
+
+
+@pytest.mark.parametrize(
+    "dtype", ["float64", "float32", "complex128", "complex64"]
+)
+def test_qr_stabilized_sign_bug(dtype):
+    from quimb.tensor.decomp import qr_stabilized
+
+    for _ in range(10):
+        Q = qu.rand_uni(4, dtype=dtype)
+        R = qu.rand_matrix(4, dtype=dtype)
+
+        # make R strictly upper triangular
+        ii, jj = np.indices(R.shape)
+        R[ii >= jj] = 0.0
+
+        X = Q @ R
+        Q2, _, R2 = qr_stabilized(X)
+
+        assert abs(np.linalg.norm((Q2 @ R2) - X)) < (
+            1e-12 if dtype in ("float64", "complex128") else 1e-6
+        )

trajectree/quimb/tests/test_tensor/test_mera.py

@@ -0,0 +1,52 @@
+import pytest
+import quimb as qu
+import quimb.tensor as qt
+
+
+class TestMERA:
+    @pytest.mark.parametrize("dtype", [float, complex])
+    def test_construct_random(self, dtype):
+        mera = qt.MERA.rand(16)
+
+        # test outer inds
+        assert set(mera.outer_inds()) == {f"k{i}" for i in range(16)}
+
+        # test normalized
+        assert (mera.H & mera) ^ all == pytest.approx(1.0)
+
+        # test auto contract all
+        assert mera.H @ mera == pytest.approx(1.0)
+
+        # test dense conversion
+        md = mera.to_qarray()
+        assert md.H @ md == pytest.approx(1.0)
+
+    def test_1d_vector_methods(self):
+        X = qu.spin_operator("X", sparse=True)
+
+        mera = qt.MERA.rand(16)
+        meraX = mera.gate(X.toarray(), 7)
+        assert mera is not meraX
+        x1 = mera.H @ meraX
+
+        md = mera.to_qarray()
+        mdX = qu.ikron(X, [2] * 16, 7) @ md
+        x2 = md.H @ mdX
+        # check against dense
+        assert x1 == pytest.approx(x2)
+
+        # check 'outside lightcone' unaffected
+        assert mera.select(3).H @ meraX.select(3) == pytest.approx(1.0)
+
+        # check only need 'lightcone' to compute local
+        assert mera.select(7).H @ meraX.select(7) == pytest.approx(x2)
+
+    @pytest.mark.parametrize("method", ["qr", "exp", "cayley", "mgs", "svd"])
+    def test_isometrize(self, method):
+        mera = qt.MERA.rand(16, dangle=True)
+        assert mera.H @ mera == pytest.approx(2.0)
+        for t in mera:
+            t.modify(data=qu.randn(t.shape))
+        assert mera.H @ mera != pytest.approx(2.0)
+        mera.isometrize_(method=method)
+        assert mera.H @ mera == pytest.approx(2.0)

trajectree/quimb/tests/test_tensor/test_optimizers.py

@@ -0,0 +1,488 @@
+import functools
+import importlib
+
+import pytest
+import numpy as np
+from numpy.testing import assert_allclose
+from autoray import real
+
+import quimb as qu
+import quimb.tensor as qtn
+from quimb.tensor.optimize import Vectorizer, parse_network_to_backend
+
+
+found_torch = importlib.util.find_spec("torch") is not None
+found_autograd = importlib.util.find_spec("autograd") is not None
+found_jax = importlib.util.find_spec("jax") is not None
+found_tensorflow = importlib.util.find_spec("tensorflow") is not None
+
+if found_tensorflow:
+    import tensorflow.experimental.numpy as tnp
+
+    tnp.experimental_enable_numpy_behavior()
+
+jax_case = pytest.param(
+    "jax", marks=pytest.mark.skipif(not found_jax, reason="jax not installed")
+)
+autograd_case = pytest.param(
+    "autograd",
+    marks=pytest.mark.skipif(
+        not found_autograd, reason="autograd not installed"
+    ),
+)
+tensorflow_case = pytest.param(
+    "tensorflow",
+    marks=pytest.mark.skipif(
+        not found_tensorflow, reason="tensorflow not installed"
+    ),
+)
+pytorch_case = pytest.param(
+    "torch",
+    marks=pytest.mark.skipif(not found_torch, reason="pytorch not installed"),
+)
+
+
+@pytest.fixture
+def tagged_qaoa_tn():
+    """
+    make qaoa tensor network, with RZZ and RX tagged on a per-round basis
+    so that these tags can be used as shared_tags to TNOptimizer
+    """
+
+    n = 8
+    depth = 4
+    terms = [(i, (i + 1) % n) for i in range(n)]
+    gammas = qu.randn(depth)
+    betas = qu.randn(depth)
+
+    # make circuit
+    circuit_opts = {"gate_opts": {"contract": False}}
+    circ = qtn.Circuit(n, **circuit_opts)
+
+    # layer of hadamards to get into plus state
+    for i in range(n):
+        circ.apply_gate("H", i, gate_round=0)
+
+    for d in range(depth):
+        for i, j in terms:
+            circ.apply_gate(
+                "RZZ", -gammas[d], i, j, gate_round=d, parametrize=True
+            )
+
+        for i in range(n):
+            circ.apply_gate(
+                "RX", betas[d] * 2, i, gate_round=d, parametrize=True
+            )
+
+    # tag circuit for shared_tags
+    tn_tagged = circ.psi.copy()
+    for i in range(depth):
+        tn_tagged.select(["RZZ", f"ROUND_{i}"]).add_tag(f"p{2 * i}")
+        tn_tagged.select(["RX", f"ROUND_{i}"]).add_tag(f"p{2 * i + 1}")
+
+    return n, depth, tn_tagged
+
+
+@pytest.fixture
+def heis_pbc():
+    L = 10
+    chi = 8
+    dtype = "float32"
+    psi0 = qtn.MPS_rand_state(L, chi, cyclic=True, seed=42).astype(dtype)
+    H = qtn.MPO_ham_heis(L, cyclic=True).astype(dtype)
+
+    def norm_fn(psi):
+        factor = (psi & psi).contract(all, optimize="auto-hq")
+        return psi / factor**0.5
+
+    def loss_fn(psi, H):
+        k, H, b = qtn.tensor_network_align(psi, H, psi)
+        energy = (k & H & b).contract(all, optimize="auto-hq")
+        return energy
+
+    en_ex = qu.groundenergy(qu.ham_heis(L, cyclic=True, sparse=True))
+
+    return psi0, H, norm_fn, loss_fn, en_ex
+
+
+@pytest.fixture
+def ham_mbl_pbc_complex():
+    L = 10
+    chi = 8
+    dtype = "complex64"
+    psi0 = qtn.MPS_rand_state(L, chi, cyclic=True, seed=42).astype(dtype)
+
+    ham_opts = {"cyclic": True, "dh": 0.7, "dh_dim": 3, "seed": 42}
+    H = qtn.MPO_ham_mbl(L, **ham_opts).astype(dtype)
+
+    def norm_fn(psi):
+        factor = (psi.H & psi).contract(all, optimize="auto-hq")
+        return psi * factor**-0.5
+
+    def loss_fn(psi, H):
+        k, H, b = qtn.tensor_network_align(psi, H, psi.H)
+        energy = (k & H & b).contract(all, optimize="auto-hq")
+        return real(energy)
+
+    en_ex = qu.groundenergy(qu.ham_mbl(L, sparse=True, **ham_opts))
+
+    return psi0, H, norm_fn, loss_fn, en_ex
+
+
+def test_vectorizer():
+    shapes = [(2, 3), (4, 5), (6, 7, 8)]
+    dtypes = ["complex64", "float32", "complex64"]
+    arrays = [qu.randn(s, dtype=dtype) for s, dtype in zip(shapes, dtypes)]
+
+    v = Vectorizer(arrays)
+
+    grads = [qu.randn(s, dtype=dtype) for s, dtype in zip(shapes, dtypes)]
+    v.pack(grads, "grad")
+
+    new_arrays = v.unpack(v.vector)
+    for x, y in zip(arrays, new_arrays):
+        assert_allclose(x, y)
+
+    new_arrays = v.unpack(v.grad)
+    for x, y in zip(grads, new_arrays):
+        assert_allclose(x, y)
+
+
+def rand_array(rng):
+    ndim = rng.integers(1, 6)
+    shape = rng.integers(2, 6, size=ndim)
+    dtype = rng.choice(["float32", "float64", "complex64", "complex128"])
+    x = rng.normal(shape).astype(dtype)
+    if "complex" in dtype:
+        x += 1j * rng.normal(shape).astype(dtype)
+    return x
+
+
+def random_array_pytree(rng, max_depth=3):
+    def _recurse(d=0):
+        if d >= max_depth:
+            return rand_array(rng)
+        t = rng.choice(["array", "list", "tuple", "dict"])
+        if t == "array":
+            return rand_array(rng)
+        elif t == "list":
+            return [_recurse(d + 1) for _ in range(rng.integers(2, 6))]
+        elif t == "tuple":
+            return tuple([_recurse(d + 1) for _ in range(rng.integers(2, 6))])
+        elif t == "dict":
+            cs = (chr(i) for i in range(ord("a"), ord("z") + 1))
+            return {
+                next(cs): _recurse(d + 1) for _ in range(rng.integers(2, 6))
+            }
+
+    return _recurse()
+
+
+def test_vectorizer_pytree():
+    tree = random_array_pytree(np.random.default_rng(666))
+    v = Vectorizer(tree)
+    new_tree = v.unpack()
+    assert tree is not new_tree
+    assert str(tree) == str(new_tree)
+
+
+@pytest.mark.parametrize(
+    "backend", [jax_case, autograd_case, tensorflow_case, pytorch_case]
+)
+@pytest.mark.parametrize("method", ["simple", "basin"])
+def test_optimize_pbc_heis(heis_pbc, backend, method):
+    psi0, H, norm_fn, loss_fn, en_ex = heis_pbc
+    tnopt = qtn.TNOptimizer(
+        psi0,
+        loss_fn,
+        norm_fn,
+        loss_constants={"H": H},
+        autodiff_backend=backend,
+    )
+    if method == "simple":
+        psi_opt = tnopt.optimize(100)
+    elif method == "basin":
+        psi_opt = tnopt.optimize_basinhopping(25, 4)
+    assert loss_fn(psi_opt, H) == pytest.approx(en_ex, rel=1e-2)
+
+
+@pytest.mark.parametrize("backend", [jax_case, autograd_case, tensorflow_case])
+@pytest.mark.parametrize("method", ["simple", "basin"])
+def test_optimize_ham_mbl_complex(ham_mbl_pbc_complex, backend, method):
+    psi0, H, norm_fn, loss_fn, en_ex = ham_mbl_pbc_complex
+    tnopt = qtn.TNOptimizer(
+        psi0,
+        loss_fn,
+        norm_fn,
+        loss_constants={"H": H},
+        autodiff_backend=backend,
+    )
+    if method == "simple":
+        psi_opt = tnopt.optimize(100)
+    elif method == "basin":
+        psi_opt = tnopt.optimize_basinhopping(25, 4)
+    assert loss_fn(psi_opt, H) == pytest.approx(en_ex, rel=1e-2)
+
+
+@pytest.mark.parametrize(
+    "backend", [jax_case, autograd_case, pytorch_case, tensorflow_case]
+)
+def test_every_parametrized_gate(backend):
+    circ = qtn.Circuit(2)
+    circ.rx(*qu.randn(1), 0, parametrize=True, tags=["OPTIMIZE"])
+    circ.ry(*qu.randn(1), 1, parametrize=True, tags=["OPTIMIZE"])
+    circ.rz(*qu.randn(1), 0, parametrize=True, tags=["OPTIMIZE"])
+    circ.rxx(*qu.randn(1), 0, 1, parametrize=True, tags=["OPTIMIZE"])
+    circ.ryy(*qu.randn(1), 0, 1, parametrize=True, tags=["OPTIMIZE"])
+    circ.rzz(*qu.randn(1), 1, 0, parametrize=True, tags=["OPTIMIZE"])
+    circ.crx(*qu.randn(1), 1, 0, parametrize=True, tags=["OPTIMIZE"])
+    circ.cry(*qu.randn(1), 0, 1, parametrize=True, tags=["OPTIMIZE"])
+    circ.crz(*qu.randn(1), 1, 0, parametrize=True, tags=["OPTIMIZE"])
+    circ.u1(*qu.randn(1), 1, parametrize=True, tags=["OPTIMIZE"])
+    circ.u2(*qu.randn(2), 0, parametrize=True, tags=["OPTIMIZE"])
+    circ.u3(*qu.randn(3), 1, parametrize=True, tags=["OPTIMIZE"])
+    circ.cu1(*qu.randn(1), 0, 1, parametrize=True, tags=["OPTIMIZE"])
+    circ.cu2(*qu.randn(2), 1, 0, parametrize=True, tags=["OPTIMIZE"])
+    circ.cu3(*qu.randn(3), 1, 0, parametrize=True, tags=["OPTIMIZE"])
+    circ.fsim(*qu.randn(2), 0, 1, parametrize=True, tags=["OPTIMIZE"])
+    circ.fsimg(*qu.randn(5), 1, 0, parametrize=True, tags=["OPTIMIZE"])
+    circ.givens(*qu.randn(1), 0, 1, parametrize=True, tags=["OPTIMIZE"])
+    circ.givens2(*qu.randn(2), 0, 1, parametrize=True, tags=["OPTIMIZE"])
+    circ.su4(*qu.randn(15), 0, 1, parametrize=True, tags=["OPTIMIZE"])
+    psi = circ.psi
+
+    target = qtn.Dense1D(qu.rand_haar_state(4))
+
+    def loss(psi, target):
+        return -(real(psi.H @ target) ** 2)
+
+    f0 = loss(psi, target)
+    tnopt = qtn.TNOptimizer(
+        psi,
+        loss,
+        tags="OPTIMIZE",
+        loss_constants={"target": target},
+        autodiff_backend=backend,
+        device="cpu",
+    )
+    tnopt.optimize(1)
+    assert tnopt.loss < f0
+
+
+@pytest.mark.parametrize("backend", [jax_case, autograd_case, tensorflow_case])
+def test_parametrized_circuit(backend):
+    H = qu.ham_mbl(4, dh=3.0, dh_dim=3)
+    gs = qu.groundstate(H)
+    T_gs = qtn.Dense1D(gs)
+
+    def loss(psi, target):
+        f = psi.H & target
+        f.rank_simplify_()
+        return -abs(f ^ all)
+
+    circ = qtn.circ_ansatz_1D_brickwork(4, depth=4)
+    psi0 = circ.psi
+    tnopt = qtn.TNOptimizer(
+        psi0,
+        loss,
+        tags="U3",
+        loss_constants=dict(target=T_gs),
+        autodiff_backend=backend,
+        loss_target=-0.99,
+    )
+    psi_opt = tnopt.optimize(20)
+    assert sum(loss < -0.99 for loss in tnopt.losses) == 1
+    assert qu.fidelity(psi_opt.to_dense(), gs) > 0.99
+
+
+def mera_norm_fn(mera):
+    return mera.isometrize(method="cayley")
+
+
+def mera_local_expectation(mera, terms, where):
+    tags = [mera.site_tag(coo) for coo in where]
+    mera_ij = mera.select(tags, "any")
+    mera_ij_G = mera_ij.gate(terms[where], where)
+    mera_ij_ex = mera_ij_G & mera_ij.H
+    return mera_ij_ex.contract(all, optimize="auto-hq")
+
+
+@pytest.mark.parametrize(
+    "backend", [autograd_case, jax_case, tensorflow_case, pytorch_case]
+)
+@pytest.mark.parametrize("executor", [None, "threads"])
+def test_multiloss(backend, executor):
+    if executor == "threads":
+        from concurrent.futures import ThreadPoolExecutor
+
+        executor = ThreadPoolExecutor(2)
+
+    L = 8
+    D = 3
+    dtype = "float32"
+
+    mera = qtn.MERA.rand(L, max_bond=D, dtype=dtype)
+
+    H2 = qu.ham_heis(2).real.astype(dtype)
+    terms = {(i, (i + 1) % L): H2 for i in range(L)}
+
+    loss_fns = [
+        functools.partial(mera_local_expectation, where=where)
+        for where in terms
+    ]
+
+    tnopt = qtn.TNOptimizer(
+        mera,
+        loss_fn=loss_fns,
+        norm_fn=mera_norm_fn,
+        loss_constants={"terms": terms},
+        autodiff_backend=backend,
+        executor=executor,
+        device="cpu",
+    )
+
+    tnopt.optimize(10)
+    # ex = -3.6510934089371734
+    assert tnopt.loss < -2.5
+
+    if executor is not None:
+        executor.shutdown()
+
+
+def test_parse_network_to_backend_shared_tags(tagged_qaoa_tn):
+    n, depth, psi0 = tagged_qaoa_tn
+
+    def to_constant(x):
+        return np.asarray(x)
+
+    tags = [f"p{i}" for i in range(2 * depth)]
+    tn_tagged, variabes = parse_network_to_backend(
+        psi0,
+        tags=tags,
+        shared_tags=tags,
+        to_constant=to_constant,
+    )
+    # test number of variables identified
+    assert len(variabes) == 2 * depth
+    # each variable tag should be in n tensors
+    for i in range(len(tags)):
+        var_tag = f"__VARIABLE{i}__"
+        assert len(tn_tagged.select(var_tag).tensors) == n
+
+
+def test_parse_network_to_backend_individual_tags(tagged_qaoa_tn):
+    n, depth, psi0 = tagged_qaoa_tn
+
+    def to_constant(x):
+        return np.asarray(x)
+
+    tags = [f"p{i}" for i in range(2 * depth)]
+    tn_tagged, variabes = parse_network_to_backend(
+        psi0, tags=tags, to_constant=to_constant
+    )
+    # test number of variables identified
+    assert len(variabes) == 2 * depth * n
+    # each variable tag should only be in 1 tensors
+    for i in range(len(tags)):
+        var_tag = f"__VARIABLE{i}__"
+        assert len(tn_tagged.select_tensors(var_tag)) == 1
+
+
+def test_parse_network_to_backend_constant_tags(tagged_qaoa_tn):
+    n, depth, psi0 = tagged_qaoa_tn
+
+    def to_constant(x):
+        return np.asarray(x)
+
+    # constant tags, include shared variable tags for first QAOA layer
+    constant_tags = ["PSI0", "H", "p0", "p1"]
+    tn_tagged, variabes = parse_network_to_backend(
+        psi0, constant_tags=constant_tags, to_constant=to_constant
+    )
+
+    # test number of variables identified
+    assert len(variabes) == 2 * (depth - 1) * n
+    # each variable tag should only be in 1 tensors
+    for i in range(len(variabes)):
+        var_tag = f"__VARIABLE{i}__"
+        assert len(tn_tagged.select(var_tag).tensors) == 1
+
+
+@pytest.mark.parametrize("backend", [jax_case, autograd_case, tensorflow_case])
+def test_shared_tags(tagged_qaoa_tn, backend):
+    n, depth, psi0 = tagged_qaoa_tn
+
+    H = qu.ham_heis(
+        n,
+        j=(0.0, 0.0, -1.0),
+        b=(1.0, 0.0, 0.0),
+        cyclic=True,
+    )
+    gs = qu.groundstate(H)
+    T_gs = qtn.Dense1D(gs).astype(complex)  # tensorflow needs all same dtype
+
+    def loss(psi, target):
+        f = psi.H & target
+        f.rank_simplify_()
+        return -abs(f ^ all)
+
+    tags = [f"p{i}" for i in range(2 * depth)]
+    tnopt = qtn.TNOptimizer(
+        psi0,
+        loss_fn=loss,
+        tags=tags,
+        shared_tags=tags,
+        loss_constants={"target": T_gs},
+        autodiff_backend=backend,
+        # loss_target=-0.99,
+    )
+
+    # run optimisation and test output
+    psi_opt = tnopt.optimize_basinhopping(n=10, nhop=5)
+    # assert sum(loss < -0.99 for loss in tnopt.losses) == 1
+    assert qu.fidelity(psi_opt.to_dense(), gs) > 0.99
+
+    # test dimension of optimisation space
+    assert tnopt.res.x.size == 2 * depth
+
+    # examine tensors inside optimised TN and check sharing was done
+    for tag in tags:
+        test_data = None
+        for t in psi_opt.select_tensors(tag):
+            if test_data is None:
+                test_data = t.get_params()
+            else:
+                assert_allclose(test_data, t.get_params())
+
+
+@pytest.mark.parametrize(
+    "backend", [jax_case, autograd_case, tensorflow_case, pytorch_case]
+)
+@pytest.mark.parametrize("simplify", ["ADCRS", "R"])
+def test_optimize_circuit_directly(backend, simplify):
+    if backend == "jax" and simplify == "ADCRS":
+        pytest.skip("JAX does not support dynamic simplification.")
+
+    circ = qtn.Circuit(2)
+    rng = np.random.default_rng(42)
+    circ.u3(*rng.uniform(high=2 * np.pi, size=3), 0, parametrize=True)
+    circ.u3(*rng.uniform(high=2 * np.pi, size=3), 1, parametrize=True)
+    circ.cnot(0, 1)
+    circ.u3(*rng.uniform(high=2 * np.pi, size=3), 0, parametrize=True)
+    circ.u3(*rng.uniform(high=2 * np.pi, size=3), 1, parametrize=True)
+
+    H = qu.ham_heis(2).astype("complex128")
+
+    def loss(circ, H):
+        return real(
+            circ.local_expectation(H, (0, 1), simplify_sequence=simplify)
+        )
+
+    assert loss(circ, H) > -0.74
+    tnopt = qtn.TNOptimizer(
+        circ, loss, loss_constants=dict(H=H), autodiff_backend=backend
+    )
+    circ_opt = tnopt.optimize(10)
+    assert circ_opt is not circ
+    assert loss(circ_opt, H) < -0.74
+    assert {t.backend for t in circ_opt.psi} == {"numpy"}