tensorcircuit-nightly 1.3.0.dev20250728__py3-none-any.whl → 1.4.0.dev20251103__py3-none-any.whl

tests/test_stabilizer.py

@@ -1,226 +0,0 @@
-import os
-import sys
-
-import numpy as np
-import pytest
-
-thisfile = os.path.abspath(__file__)
-modulepath = os.path.dirname(os.path.dirname(thisfile))
-
-sys.path.insert(0, modulepath)
-import tensorcircuit as tc
-
-# Skip all tests if stim is not installed
-stim = pytest.importorskip("stim")
-
-
-def test_basic_gates():
-    c = tc.StabilizerCircuit(2)
-    c.h(0)
-    c.cnot(0, 1)
-    results = c.measure(0, 1)
-    print(results)
-    results = c.measure(0, 1)
-    print(results)
-    assert len(results) == 2
-    assert all(r in [0, 1] for r in results)
-
-
-def test_bell_state():
-    c = tc.StabilizerCircuit(2)
-    c.H(1)
-    c.cx(1, 0)
-    # Test Z measurements correlation
-    samples = c.sample(batch=1000)
-    assert samples.shape == (1000, 2)
-    counts = np.sum(samples, axis=0)
-    # Should be roughly equal number of 00 and 11 states
-    assert abs(counts[0] - counts[1]) < 50
-
-
-def test_ghz_state():
-    c = tc.StabilizerCircuit(3)
-    c.h(0)
-    c.cnot(0, 1)
-    c.cnot(1, 2)
-    # Test expectation values
-    exp_z = c.expectation_ps(z=[0, 1])
-    np.testing.assert_allclose(exp_z, 1.0, atol=1e-6)
-
-
-def test_stabilizer_operations():
-    c = tc.StabilizerCircuit(2)
-    # Test all supported gates
-    for gate in ["h", "x", "y", "z", "s", "sdg"]:
-        getattr(c, gate)(0)
-    c.cnot(0, 1)
-    c.cz(0, 1)
-    c.swap(0, 1)
-    print(c.current_circuit())
-
-
-def test_sample_expectation():
-    c = tc.StabilizerCircuit(2)
-    c.h(0)
-    c.cnot(0, 1)
-    # Test sampling-based expectation
-    exp = c.sample_expectation_ps(z=[0, 1], shots=1000)
-    assert abs(exp - 1.0) < 0.1
-
-
-def test_invalid_gates():
-    c = tc.StabilizerCircuit(1)
-    with pytest.raises(ValueError):
-        c.t(0)
-
-
-def test_qir():
-    c = tc.StabilizerCircuit(2)
-    c.h(0)
-    c.cnot(0, 1)
-    qir = c.to_qir()
-    assert len(qir) == 2
-    assert qir[0]["name"] == "h"
-    assert qir[1]["name"] == "cnot"
-    print(qir)
-    c1 = tc.Circuit.from_qir(qir)
-    np.testing.assert_allclose(c1.expectation_ps(z=(0, 1)), 1, atol=1e-5)
-
-
-def test_cond_measure():
-    c = tc.StabilizerCircuit(3)
-
-    # Prepare Bell pair between qubits 1 and 2
-    c.H(1)
-    c.CNOT(1, 2)
-
-    # Prepare state to teleport on qubit 0 (can only be Clifford state)
-    c.X(0)
-
-    # Teleportation circuit
-    c.CNOT(0, 1)
-    c.H(0)
-
-    # Measure qubits 0 and 1
-    r0 = c.cond_measure(0)
-    r1 = c.cond_measure(1)
-    # Apply corrections based on measurements
-    if r0 == 1:
-        c.Z(2)
-    if r1 == 1:
-        c.X(2)
-
-    # Verify teleported state
-    final = c.measure(2)
-    np.testing.assert_allclose(final, 1)
-
-
-def test_post_select():
-    c = tc.StabilizerCircuit(2)
-    c.h(0)
-    c.s(0)
-    c.cx(0, 1)
-    c.post_select(1, keep=1)
-    np.testing.assert_allclose(c.expectation_ps(z=[0]), -1, atol=1e-5)
-
-
-def test_to_openqasm():
-    c = tc.StabilizerCircuit(3)
-    c.sd(0)
-    c.cz(0, 1)
-    c.s(2)
-    c.measure(0, 1)
-
-    # Test basic circuit conversion
-    qasm = c.to_openqasm()
-    print(qasm)
-
-    c1 = tc.StabilizerCircuit.from_openqasm(qasm)
-    print(c1.draw())
-
-
-def test_ee():
-    c = tc.Circuit(8)
-    for i in range(3):
-        c.h(i)
-        c.cx(i, i + 4)
-        c.sd(i + 2)
-    ee0 = tc.quantum.entanglement_entropy(c.state(), list(range(4)))
-    c1 = tc.StabilizerCircuit.from_openqasm(c.to_openqasm())
-    ee1 = c1.entanglement_entropy(list(range(4)))
-    np.testing.assert_allclose(ee0, ee1, atol=1e-6)
-
-
-def test_random_gates():
-    c = tc.StabilizerCircuit(4)
-    c.random_gate(0, 1, recorded=True)
-    c.random_gate(2, 3)
-    c.random_gate(1, 2)
-    print(c.entanglement_entropy(list(range(2))))
-    print(len(c.current_circuit()))
-
-
-def test_circuit_state():
-    c = tc.StabilizerCircuit(2)
-    c.h(1)
-    c1 = tc.Circuit(2)
-    c1.h(1)
-    np.testing.assert_allclose(c.state(), c1.state(), atol=1e-5)
-
-
-def test_circuit_inputs():
-    c = tc.StabilizerCircuit(2, inputs=[stim.PauliString("XX"), stim.PauliString("ZZ")])
-    c.cnot(0, 1)
-    c.h(0)
-    np.testing.assert_allclose(c.expectation_ps(z=[0]), 1, atol=1e-6)
-    np.testing.assert_allclose(c.expectation_ps(z=[1]), 1, atol=1e-6)
-
-
-def test_depolarize():
-    r = []
-    for _ in range(40):
-        c = tc.StabilizerCircuit(2)
-        c.h(0)
-        c.depolarizing(0, 1, p=0.2)
-        c.h(0)
-        r.append(c.expectation_ps(z=[0]))
-    assert 4 < np.sum(r) < 39
-
-
-def test_tableau_inputs():
-    c = tc.StabilizerCircuit(2)
-    c.x(1)
-    c.s(1)
-    it = c.current_inverse_tableau()
-    c1 = tc.StabilizerCircuit(2, tableau_inputs=it)
-    c1.s(1)
-    c1.x(1)
-    np.testing.assert_allclose(c1.state()[0], 1, atol=1e-6)
-
-
-def test_mipt():
-    resource = [stim.Tableau.random(2) for _ in range(1000)]
-
-    def ruc(n, nlayer, p):
-        c = tc.StabilizerCircuit(n)
-        status = np.random.choice(1000, size=[n, nlayer], replace=True)
-        for j in range(nlayer):
-            for i in range(0, n, 2):
-                c.tableau_gate(i, (i + 1) % n, tableau=resource[status[i, j]])
-            for i in range(1, n, 2):
-                c.tableau_gate(i, (i + 1) % n, tableau=resource[status[i, j]])
-            mask = np.random.random(n) < p
-            ids = list(np.where(mask)[0])
-            c.cond_measure_many(*ids)
-        return c.entanglement_entropy(list(range(n // 2)))
-
-    print(ruc(50, 10, 0.1))
-
-
-def test_measure_with_prob():
-    c = tc.StabilizerCircuit(3)
-    c.h(0)
-    c.cnot(0, 1)
-    m, p = c.measure(0, 2, with_prob=True)
-    np.testing.assert_allclose(p, 0.5, atol=1e-6)
-    print(m)

tests/test_templates.py

@@ -1,218 +0,0 @@
-# pylint: disable=invalid-name
-
-import os
-import sys
-
-import numpy as np
-import pytest
-from pytest_lazyfixture import lazy_fixture as lf
-
-thisfile = os.path.abspath(__file__)
-modulepath = os.path.dirname(os.path.dirname(thisfile))
-
-sys.path.insert(0, modulepath)
-import tensorcircuit as tc
-
-
-def test_any_measurement():
-    c = tc.Circuit(2)
-    c.H(0)
-    c.H(1)
-    mea = np.array([1, 1])
-    r = tc.templates.measurements.any_measurements(c, mea, onehot=True)
-    np.testing.assert_allclose(r, 1.0, atol=1e-5)
-    mea2 = np.array([3, 0])
-    r2 = tc.templates.measurements.any_measurements(c, mea2, onehot=True)
-    np.testing.assert_allclose(r2, 0.0, atol=1e-5)
-
-
-@pytest.mark.parametrize("backend", [lf("jaxb"), lf("tfb")])
-def test_parameterized_local_measurement(backend):
-    c = tc.Circuit(3)
-    c.X(0)
-    c.cnot(0, 1)
-    c.H(-1)
-    basis = tc.backend.convert_to_tensor(np.array([3, 3, 1]))
-    r = tc.templates.measurements.parameterized_local_measurements(
-        c, structures=basis, onehot=True
-    )
-    np.testing.assert_allclose(r, np.array([-1, -1, 1]), atol=1e-5)
-
-
-@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
-def test_sparse_expectation(backend):
-    ham = tc.backend.coo_sparse_matrix(
-        indices=[[0, 1], [1, 0]], values=tc.backend.ones([2]), shape=(2, 2)
-    )
-
-    def f(param):
-        c = tc.Circuit(1)
-        c.rx(0, theta=param[0])
-        c.H(0)
-        return tc.templates.measurements.sparse_expectation(c, ham)
-
-    fvag = tc.backend.jit(tc.backend.value_and_grad(f))
-    param = tc.backend.zeros([1])
-    print(fvag(param))
-
-
-def test_bell_block():
-    c = tc.Circuit(4)
-    c = tc.templates.blocks.Bell_pair_block(c)
-    for _ in range(10):
-        s = c.perfect_sampling()[0]
-        assert s[0] != s[1]
-        assert s[2] != s[3]
-
-
-def test_qft_block() -> None:
-    n_qubits = 4
-    c = tc.Circuit(n_qubits)
-    c = tc.templates.blocks.qft(c, *range(n_qubits))
-    mat = c.quoperator().eval().reshape(2 ** (n_qubits), -1)
-    N = 2**n_qubits
-    ref = np.exp(
-        1j * 2 * np.pi * np.arange(N).reshape(-1, 1) * np.arange(N).reshape(1, -1) / N
-    ) / np.sqrt(N)
-    np.testing.assert_allclose(mat, ref, atol=1e-7)
-
-    c = tc.Circuit(n_qubits)
-    c = tc.templates.blocks.qft(c, *range(n_qubits), inverse=True)
-    mat = c.quoperator().eval().reshape(2 ** (n_qubits), -1)
-    np.testing.assert_allclose(mat, ref.T.conj(), atol=1e-7)
-
-
-def test_grid_coord():
-    cd = tc.templates.graphs.Grid2DCoord(3, 2)
-    assert cd.all_cols() == [(0, 3), (1, 4), (2, 5)]
-    assert cd.all_rows() == [(0, 1), (1, 2), (3, 4), (4, 5)]
-
-
-@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
-def test_qaoa_template(backend):
-    cd = tc.templates.graphs.Grid2DCoord(3, 2)
-    g = cd.lattice_graph(pbc=False)
-    for e1, e2 in g.edges:
-        g[e1][e2]["weight"] = np.random.uniform()
-
-    def forward(paramzz, paramx):
-        c = tc.Circuit(6)
-        for i in range(6):
-            c.H(i)
-        c = tc.templates.blocks.QAOA_block(c, g, paramzz, paramx)
-        return tc.templates.measurements.spin_glass_measurements(c, g)
-
-    fvag = tc.backend.jit(tc.backend.value_and_grad(forward, argnums=(0, 1)))
-    paramzz = tc.backend.real(tc.backend.ones([1]))
-    paramx = tc.backend.real(tc.backend.ones([1]))
-    _, gr = fvag(paramzz, paramx)
-    np.testing.assert_allclose(gr[1].shape, [1])
-    paramzz = tc.backend.real(tc.backend.ones([7]))
-    paramx = tc.backend.real(tc.backend.ones([1]))
-    _, gr = fvag(paramzz, paramx)
-    np.testing.assert_allclose(gr[0].shape, [7])
-    paramzz = tc.backend.real(tc.backend.ones([1]))
-    paramx = tc.backend.real(tc.backend.ones([6]))
-    _, gr = fvag(paramzz, paramx)
-    np.testing.assert_allclose(gr[0].shape, [1])
-    np.testing.assert_allclose(gr[1].shape, [6])
-
-
-def test_state_wrapper():
-    Bell_pair_block_state = tc.templates.blocks.state_centric(
-        tc.templates.blocks.Bell_pair_block
-    )
-    s = Bell_pair_block_state(np.array([1.0, 0, 0, 0]))
-    np.testing.assert_allclose(
-        s, np.array([0.0, 0.70710677 + 0.0j, -0.70710677 + 0.0j, 0]), atol=1e-5
-    )
-
-
-@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
-def test_amplitude_encoding(backend):
-    batched_amplitude_encoding = tc.backend.vmap(
-        tc.templates.dataset.amplitude_encoding, vectorized_argnums=0
-    )
-    figs = np.stack([np.eye(2), np.ones([2, 2])])
-    figs = tc.array_to_tensor(figs)
-    states = batched_amplitude_encoding(figs, 3)
-    # note that you cannot use nqubits=3 here for jax backend
-    # see this issue: https://github.com/google/jax/issues/7465
-    np.testing.assert_allclose(states[1], np.array([0.5, 0.5, 0.5, 0.5, 0, 0, 0, 0]))
-    states = batched_amplitude_encoding(
-        figs, 2, tc.array_to_tensor(np.array([0, 3, 1, 2]), dtype="int32")
-    )
-    np.testing.assert_allclose(states[0], 1 / np.sqrt(2) * np.array([1, 1, 0, 0]))
-
-
-@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
-def test_mpo_measurement(backend):
-    def f(theta):
-        mpo = tc.quantum.QuOperator.from_local_tensor(
-            tc.array_to_tensor(tc.gates._x_matrix), [2, 2, 2], [0]
-        )
-        c = tc.Circuit(3)
-        c.ry(0, theta=theta)
-        c.H(1)
-        c.H(2)
-        e = tc.templates.measurements.mpo_expectation(c, mpo)
-        return e
-
-    v, g = tc.backend.jit(tc.backend.value_and_grad(f))(tc.backend.ones([]))
-
-    np.testing.assert_allclose(v, 0.84147, atol=1e-4)
-    np.testing.assert_allclose(g, 0.54032, atol=1e-4)
-
-
-@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
-def test_operator_measurement(backend):
-    mpo = tc.quantum.QuOperator.from_local_tensor(
-        tc.array_to_tensor(tc.gates._x_matrix), [2, 2], [0]
-    )
-    dense = tc.array_to_tensor(np.kron(tc.gates._x_matrix, np.eye(2)))
-    sparse = tc.quantum.PauliString2COO([1, 0])
-
-    for h in [dense, sparse, mpo]:
-
-        def f(theta):
-            c = tc.Circuit(2)
-            c.ry(0, theta=theta)
-            c.H(1)
-            e = tc.templates.measurements.operator_expectation(c, h)
-            return e
-
-        v, g = tc.backend.jit(tc.backend.value_and_grad(f))(tc.backend.ones([]))
-
-        np.testing.assert_allclose(v, 0.84147, atol=1e-4)
-        np.testing.assert_allclose(g, 0.54032, atol=1e-4)
-
-
-@pytest.fixture
-def symmetric_matrix():
-    matrix = np.array([[-5.0, -2.0], [-2.0, 6.0]])
-    nsym_matrix = np.array(
-        [[1.0, 2.0, 3.0], [2.0, 4.0, 5.0], [3.0, 5.0, 6.0], [8.0, 7.0, 6.0]]
-    )
-    return matrix, nsym_matrix
-
-
-def test_QUBO_to_Ising(symmetric_matrix):
-    matrix1, matrix2 = symmetric_matrix
-    pauli_terms, weights, offset = tc.templates.conversions.QUBO_to_Ising(matrix1)
-    n = matrix1.shape[0]
-    expected_num_terms = n + n * (n - 1) // 2
-    assert len(pauli_terms) == expected_num_terms
-    assert len(weights) == expected_num_terms
-    assert isinstance(pauli_terms, list)
-    assert isinstance(weights, np.ndarray)
-    assert isinstance(offset, float)
-    assert pauli_terms == [
-        [1, 0],
-        [0, 1],
-        [1, 1],
-    ]
-    assert all(weights == np.array([3.5, -2.0, -1.0]))
-    assert offset == -0.5
-
-    with pytest.raises(ValueError):
-        tc.templates.conversions.QUBO_to_Ising(matrix2)

tests/test_torchnn.py

@@ -1,99 +0,0 @@
-import os
-import sys
-import numpy as np
-import pytest
-from pytest_lazyfixture import lazy_fixture as lf
-
-thisfile = os.path.abspath(__file__)
-modulepath = os.path.dirname(os.path.dirname(thisfile))
-
-sys.path.insert(0, modulepath)
-
-import tensorcircuit as tc
-
-try:
-    import torch
-except ImportError:
-    pytest.skip("torch is not installed", allow_module_level=True)
-
-
-@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb"), lf("torchb")])
-def test_quantumnet(backend):
-    n = 6
-    nlayers = 2
-
-    def qpred(x, weights):
-        c = tc.Circuit(n)
-        for i in range(n):
-            c.rx(i, theta=x[i])
-        for j in range(nlayers):
-            for i in range(n - 1):
-                c.cnot(i, i + 1)
-            for i in range(n):
-                c.rx(i, theta=weights[2 * j, i])
-                c.ry(i, theta=weights[2 * j + 1, i])
-        ypred = tc.backend.stack([c.expectation_ps(x=[i]) for i in range(n)])
-        return tc.backend.real(ypred)
-
-    if tc.backend.name == "pytorch":
-        use_interface = False
-    else:
-        use_interface = True
-
-    ql = tc.TorchLayer(
-        qpred, weights_shape=[2 * nlayers, n], use_interface=use_interface
-    )
-
-    yp = ql(torch.ones([3, n]))
-    print(yp)
-
-    np.testing.assert_allclose(yp.shape, np.array([3, n]))
-
-
-@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
-def test_inputs_multiple(backend):
-    n = 3
-    p = 0.1
-    K = tc.backend
-    torchb = tc.get_backend("pytorch")
-
-    def f(state, noise, weights):
-        c = tc.Circuit(n, inputs=state)
-        for i in range(n):
-            c.rz(i, theta=weights[i])
-        for i in range(n):
-            c.depolarizing(i, px=p, py=p, pz=p, status=noise[i])
-        return K.real(c.expectation_ps(x=[0]))
-
-    layer = tc.TorchLayer(f, [n], use_vmap=True, vectorized_argnums=[0, 1])
-    state = torchb.ones([2, 2**n]) / 2 ** (n / 2)
-    noise = 0.2 * torchb.ones([2, n], dtype="float32")
-    l = layer(state, noise)
-    lsum = torchb.sum(l)
-    print(l)
-    lsum.backward()
-    for p in layer.parameters():
-        print(p.grad)
-
-
-@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb"), lf("torchb")])
-def test_torchnn_hardware(backend):
-    n = 2
-
-    def qf(inputs, param):
-        inputs = tc.backend.convert_to_tensor(tc.get_backend("pytorch").numpy(inputs))
-        param = tc.backend.convert_to_tensor(tc.get_backend("pytorch").numpy(param))
-
-        c = tc.Circuit(n)
-        c.rx(0, theta=inputs[0])
-        c.rx(1, theta=inputs[1])
-        c.h(1)
-        c.rzz(0, 1, theta=param[0])
-        r = tc.backend.stack([c.expectation_ps(z=[i]) for i in range(n)])
-
-        r = tc.get_backend("pytorch").convert_to_tensor(tc.backend.numpy(r))
-        return torch.real(r)
-
-    ql = tc.torchnn.HardwareNet(qf, [1])
-    qnet = torch.nn.Sequential(ql, torch.nn.Linear(2, 1))
-    print(qnet(torch.ones([5, 2])))

tests/test_van.py

@@ -1,102 +0,0 @@
-import sys
-import os
-import itertools
-
-thisfile = os.path.abspath(__file__)
-modulepath = os.path.dirname(os.path.dirname(thisfile))
-
-sys.path.insert(0, modulepath)
-
-import numpy as np
-import tensorflow as tf
-
-
-from tensorcircuit.applications.van import MaskedLinear, ResidualBlock, MADE, PixelCNN
-
-
-def test_masklinear():
-    ml = MaskedLinear(5, 10, 3, mask=tf.zeros([10, 3, 5, 3]))
-    ml.set_weights([tf.ones([10, 3, 5, 3]), tf.ones([10, 3])])
-    tf.debugging.assert_near(ml(tf.ones([1, 5, 3])), tf.ones([1, 10, 3]))
-
-    mask = np.zeros([10, 3, 5, 3])
-    mask[3, 2, 1, 0] = 1.0
-    ml = MaskedLinear(5, 10, 3, mask=tf.constant(mask, dtype=tf.float32))
-    ml.set_weights([tf.ones([10, 3, 5, 3]), tf.zeros([10, 3])])
-    assert tf.reduce_sum(ml(tf.ones([5, 3]))[3, :]) == 1.0
-
-    w = tf.random.uniform(shape=[10, 3, 5, 3])
-    b = tf.random.uniform(shape=[10, 3])
-    w_m = tf.reshape(w, [30, 15])
-    b_m = tf.reshape(b, [30, 1])
-    inputs = tf.ones([5, 3])
-    inputs_m = tf.reshape(inputs, [15, 1])
-    r_m = w_m @ inputs_m + b_m
-    r = tf.reshape(r_m, [10, 3])
-    ml = MaskedLinear(5, 10, 3)
-    ml.set_weights([w, b])
-    tf.debugging.assert_near(ml(inputs), r)
-
-
-def test_residual_block():
-    dense1 = tf.keras.layers.Dense(10, use_bias=False)
-    dense1.build([1, 1])
-    dense1.set_weights([np.ones([1, 10])])
-    dense2 = tf.keras.layers.Dense(1, use_bias=False)
-    dense2.build([1, 10])
-    dense2.set_weights([np.ones([10, 1])])
-    m = ResidualBlock([dense1, dense2])
-    assert m(tf.ones([1, 1])) == 11.0
-
-
-def test_made():
-    import itertools
-
-    m = MADE(2, 2, 6, 3, 3, nonmerge=False)
-    l = []
-    for i in itertools.product(*[list(range(3)) for _ in range(2)]):
-        l.append(list(i))
-    basis = tf.constant(l, dtype=tf.int32)
-    print(basis)
-    ptot = tf.reduce_sum(tf.exp(m.log_prob(tf.one_hot(basis, depth=3))))
-    np.testing.assert_allclose(ptot.numpy(), 1.0, atol=1e-5)
-
-    s, logp = m.sample(10)
-    print(logp)
-    assert s.shape == (10, 2, 3)
-
-
-def test_made_fit_peak():
-    opt = tf.optimizers.Adam(learning_rate=0.01)
-    m = MADE(5, 5, 4, 3, 2)
-    for step in range(100):
-        with tf.GradientTape() as t:
-            loss = -tf.reduce_sum(
-                m.log_prob(
-                    tf.one_hot(
-                        [[0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [2, 2, 1, 1, 0]], depth=3
-                    )
-                )
-            )
-            gr = t.gradient(loss, m.variables)
-        if step % 20 == 0:
-            print(
-                tf.exp(
-                    m.log_prob(
-                        tf.one_hot(
-                            [[0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [2, 2, 1, 1, 0]], depth=3
-                        )
-                    )
-                ).numpy()
-            )
-        opt.apply_gradients(zip(gr, m.variables))
-
-
-def test_pixelcnn():
-    m = PixelCNN(3, 5, 8)
-    l = []
-    for i in itertools.product(*[list(range(3)) for _ in range(4)]):
-        l.append(list(i))
-    basis = tf.constant(tf.reshape(l, [-1, 2, 2]), dtype=tf.int32)
-    ptot = tf.reduce_sum(tf.exp(m.log_prob(tf.one_hot(basis, depth=3))))
-    np.testing.assert_allclose(ptot.numpy(), 1.0, atol=1e-5)