Recommended way of calculating the expectation value at the end of a circuit

[BrunoLiegiBastonLiegi]

What is your issue?

I am trying to train a circuit to minimize the expectation value of a given observable. I checked the documentation and the other issues I found concerning expectation values calculation, but I did not find exactly what I was looking for (sorry about the duplication in case I missed it).

Put in simple words, say that I have an hamiltonian:

$$H = - Z_0\otimes I_1 + 2 X_0 \otimes Y_1 + 3 I_0 \otimes X_1 \otimes Z_2$$

the first question is: what is the best way to represent this in quimb? At the moment I am using quimb.pauli which directly builds the corresponding array for each term, but I saw that more high level hamiltonian objects are available as well, and depending on the way you calculate the expectation values they may be more suitable, I believe.
This leads to my second and main question, supposing that I have a quimb circuit: what is the best way to calculate the expectation value of $H$ at the end of it? I am currently using quimb_circuit.local_expectation passing the corresponding pauli operator built with quimb.pauli and the sites (qubits) it is acting on. However, this means that I have to loop over each term of the hamiltonian, build the operator and call the local_expectation, whereas I've seen in some examples (e.g. this) that in some cases it could be done more simply (and efficiently I guess).

My current implementation relying on quimb_circuit.local_expectation is also having some problems when running with jax and trying to @jit or when running with pytorch.

def expectation(parameters, operators_list, sites_list, coeffs_list,):
    # build the quimb circuit using the parameters
    quimb_circuit = build_circuit(parameters)
    
    expectation_value = 0.0
    for ops, sites, coeff in zip(operators_list, sites_list, coeffs_list):
        # build quimb representation of the operator
        ops = build_quimb_operator(ops)
        # calculate the term exp val
        exp_val = quimb_circuit.local_expectation(
            ops,
            where=sites,
            backend="jax", # or "torch"
            optimize=optimizer, # I am trying both with "auto-hq" and cotengra
        )
        expectation_value += coeff * exp_val
    return expectation_value

As I said this runs fine with jax (even though it is a bit slow) but fails with:

NonConcreteBooleanIndexError: Array boolean indices must be concrete; got ShapedArray(bool[2,2])

if I try to @partial(jax.jit, static_argnums=(1,2,3)), whereas it just fails under torch with:

TypeError: tensordot(): argument 'input' (position 1) must be Tensor, not numpy.ndarray