MPS usage with NetworkState in cuQuantum python produces integer overflow

I am using an MPS simulator in the Python version of cuQuantum to compute expectation values. After defining a NetworkState with a small number of qubits and a small bond dimension (max_extent), and after defining a simple NetworkOperator to measure, the code crashes with ValueError: extents must be non-negative. This crash with small number of qubits, small bond dimension, and simple observables is not expected, and the expectation value should be feasible to compute. The code producing the crash can be found here:

https://github.com/NVIDIA/cuQuantum/blob/9da0df2fbcafb56d9e1b02998a40d3f11dd402b2/python/cuquantum/tensornet/experimental/network_state.py#L835-L840

The root cause of this issue is that the bounding of the maximum extent happens after the cumulative product, allowing the cumulative product to grow without any restriction. That results on an integer overflow even when computing, for example, the expectation value of a simple Hamiltonian for a 64-node qubit system, which should be completely feasible using an MPS representation of the state.

The following minimal code gives a way of reproducing the crash. Note that it only applies an H gate to qubit 0, and the observable is a single Pauli Z on such qubit (a simple circuit and a simple observable):

import numpy as np

from cuquantum.tensornet.experimental import (
    NetworkState,
    NetworkOperator,
    MPSConfig,
)


def minimal_repro(n=100, q=0, max_extent=64):
    dtype = np.complex128

    H = np.array(
        [[1, 1],
         [1, -1]],
        dtype=dtype,
    ) / np.sqrt(2)

    Z = np.array(
        [[1, 0],
         [0, -1]],
        dtype=dtype,
    )

    state = NetworkState(
        [2] * n,
        dtype="complex128",
        config=MPSConfig(max_extent=max_extent),
    )

    try:
        state.apply_tensor_operator(
            (q,),
            H,
            unitary=True,
        )

        op = NetworkOperator([2] * n, dtype="complex128")
        op.append_product(
            1.0,
            ((q,),),
            (Z,),
        )

        value = state.compute_expectation(
            op,
            release_workspace=True,
        )

        print(f"SUCCESS: n={n}, q={q}, value={value}")

    finally:
        state.free()


if __name__ == "__main__":
    minimal_repro(n=64, q=0, max_extent=2)

The crash only happens for values of n greater or equal than 64, reinforcing the integer overflow idea (which makes the resulting number really big and negative).

I have solved this bug locally via the following code:

cap = EXACT_MPS_EXTENT_LIMIT
if max_extent is not None:
    cap = min(cap, int(max_extent))

def bounded_cumprod(extents, cap):
    values = []
    product = 1
    for extent in extents:
        product *= int(extent)
        if product > cap:
            product = cap
        values.append(product)
    return values
combined_extents_left = bounded_cumprod(self.state_mode_extents[:-1], cap)
combined_extents_right = bounded_cumprod(self.state_mode_extents[::-1], cap)[::-1][1:]

which works and returns expectation values matching other libraries

Start at python/cuquantum/tensornet/experimental/network_state.py around lines 835-840 and run the minimal reproduction from the issue with n=64 and max_extent=2. Trace how extents are combined before the expectation calculation; done means the reproduction completes without the negative-extents ValueError and returns an expectation value matching other libraries.