Error: both inputs should have same dtype in MPSBackend._run_from_sequence_data

[v-vitale]

Dear all,

When using SequenceData and MPSBackend._run_from_sequence_data I encounter a type error, that is solved when casting all inputs to complex128.

This is a minimal example where I create a sequence, and get SequenceData from it and feed it back into MPSBackend._run_from_sequence_data.

from pulser import Register, Pulse,InterpolatedWaveform, Sequence, MockDevice
import emu_base
from emu_mps import MPSConfig,MPSBackend
import torch

reg=Register.square(2)
seq = Sequence(reg, MockDevice)
adiabatic_pulse = Pulse(
    InterpolatedWaveform(100, [0, 1, 0]),
    InterpolatedWaveform(100, [-1, 1]),         # δ ramp
    0.0,
)
seq.declare_channel("ising_global", "rydberg_global")
seq.add(adiabatic_pulse, "ising_global")

dt=10

config=MPSConfig(dt=dt, observables=[])
cd_config = config.with_changes(observables=[])

pulser_data = emu_base.pulser_adapter.PulserData(
    sequence=seq, config=cd_config, dt=cd_config.dt
)
seq0 = next(pulser_data.get_sequences())

omegas = seq0.omega.to(dtype=torch.float64)
deltas = seq0.delta.to(dtype=torch.float64)
phis = seq0.phi.to(dtype=torch.float64)
interact_mat = seq0.interaction_matrix(0.0)  # matrix is constant in time

target_times = [x * dt for x in range(0, omegas.shape[0] + 1)]
    

data=emu_base.SequenceData(omegas,
            deltas,
            phis,
            lambda x: interact_mat,
            seq.register.qubit_ids,
            bad_atoms=[False] * 2,
            lindblad_ops=[],
            state_prep_error=0.0,
            target_times=target_times,
            eigenstates=("r", "g"),
            hamiltonian_type=emu_base.HamiltonianType.Rydberg,
        )


MPSBackend._run_from_sequence_data(data, config)

Error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[13], line 51
     34 target_times = [x * dt for x in range(0, omegas.shape[0] + 1)]
     37 data=emu_base.SequenceData(omegas,
     38             deltas,
     39             phis,
   (...)
     47             hamiltonian_type=emu_base.HamiltonianType.Rydberg,
     48         )
---> 51 MPSBackend._run_from_sequence_data(data, config)

File ~/pulse_counter_diabatic/.venv/lib/python3.10/site-packages/emu_mps/mps_backend.py:74, in MPSBackend._run_from_sequence_data(sequence_data, config)
     69 @staticmethod
     70 def _run_from_sequence_data(
     71     sequence_data: SequenceData, config: MPSConfig
     72 ) -> Results:
     73     impl = create_impl(sequence_data, config)
---> 74     impl.init()  # This is separate from the constructor for testing purposes.
     75     result = MPSBackend._run(impl)
     76     return impl.permute_results(result, config.optimize_qubit_ordering)

File ~/pulse_counter_diabatic/.venv/lib/python3.10/site-packages/emu_mps/mps_backend_impl.py:323, in MPSBackendImpl.init(self)
    321 self.init_dark_qubits()
    322 self.init_initial_state(self.config.initial_state)
--> 323 self.init_noiseless_hamiltonian()
    324 self.fill_results()  # at t == 0 for pulser compatibility
    325 self.update_H()

File ~/pulse_counter_diabatic/.venv/lib/python3.10/site-packages/emu_mps/mps_backend_impl.py:287, in MPSBackendImpl.init_noiseless_hamiltonian(self)
    280 self.current_interaction_matrix = self._get_interaction_matrix()
    281 self.hamiltonian = make_H(
    282     interaction_matrix=self.current_interaction_matrix,
    283     hamiltonian_type=self.hamiltonian_type,
    284     num_gpus_to_use=self.resolved_num_gpus,
    285     dim=self.dim,
    286 )
--> 287 self.update_H_no_noise()

File ~/pulse_counter_diabatic/.venv/lib/python3.10/site-packages/emu_mps/mps_backend_impl.py:299, in MPSBackendImpl.update_H_no_noise(self)
    298 def update_H_no_noise(self) -> None:
--> 299     update_H(
    300         hamiltonian=self.hamiltonian,
    301         omega=self.omega[self._timestep_index, :],
    302         delta=self.delta[self._timestep_index, :],
    303         phi=self.phi[self._timestep_index, :],
    304         noise=torch.zeros(self.dim, self.dim, dtype=dtype),  # no noise
    305     )

File ~/pulse_counter_diabatic/.venv/lib/python3.10/site-packages/emu_mps/hamiltonian.py:498, in update_H(hamiltonian, omega, delta, phi, noise)
    493     raise ValueError(
    494         f"noise must have shape (2, 2) or (3, 3), got {tuple(noise.shape)}"
    495     )
    496 nqubits = omega.size(dim=0)
--> 498 a = torch.tensordot(omega * torch.cos(phi), Operators.sx, dims=0)
    499 c = torch.tensordot(delta, Operators.n, dims=0)
    500 b = torch.tensordot(omega * torch.sin(phi), Operators.sy, dims=0)

File ~/pulse_counter_diabatic/.venv/lib/python3.10/site-packages/torch/functional.py:1382, in tensordot(a, b, dims, out)
   1379     dims_b = list(range(dims))
   1381 if out is None:
-> 1382     return _VF.tensordot(a, b, dims_a, dims_b)  # type: ignore[attr-defined]
   1383 else:
   1384     return _VF.tensordot(a, b, dims_a, dims_b, out=out)

RuntimeError: both inputs should have same dtype

The problem doesn't arise through the standard way of providing a sequence

from pulser import Register, Pulse,InterpolatedWaveform, Sequence, MockDevice
import emu_base
from emu_mps import MPSConfig,MPSBackend
import torch

reg=Register.square(2)
seq = Sequence(reg, MockDevice)
adiabatic_pulse = Pulse(
    InterpolatedWaveform(100, [0, 1, 0]),
    InterpolatedWaveform(100, [-1, 1]),         # δ ramp
    0.0,
)
seq.declare_channel("ising_global", "rydberg_global")
seq.add(adiabatic_pulse, "ising_global")

dt=10

config=MPSConfig(dt=dt, observables=[])

backend = MPSBackend(seq, config=config)
results = backend.run()

and is fixed by casting

data=emu_base.SequenceData(omegas.to(dtype=torch.complex128),
            deltas.to(dtype=torch.complex128),
            phis.to(dtype=torch.complex128),
            lambda x: interact_mat,
            seq.register.qubit_ids,
            bad_atoms=[False] * 2,
            lindblad_ops=[],
            state_prep_error=0.0,
            target_times=target_times,
            eigenstates=("r", "g"),
            hamiltonian_type=emu_base.HamiltonianType.Rydberg,
        )