emu-mps 1.2.3__py3-none-any.whl → 1.2.5__py3-none-any.whl

emu_mps/__init__.py

@@ -10,10 +10,10 @@ from emu_base import (
     StateResult,
     SecondMomentOfEnergy,
 )
+from .mps_config import MPSConfig
 from .mpo import MPO
 from .mps import MPS, inner
 from .mps_backend import MPSBackend
-from .mps_config import MPSConfig
 
 
 __all__ = [
@@ -35,4 +35,4 @@ __all__ = [
     "SecondMomentOfEnergy",
 ]
 
-__version__ = "1.2.3"
+__version__ = "1.2.5"

emu_mps/algebra.py

@@ -1,13 +1,17 @@
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
 import math
+
+from emu_mps import MPSConfig
 from emu_mps.utils import truncate_impl
 
 
 def add_factors(
-    left: list[torch.tensor], right: list[torch.tensor]
-) -> list[torch.tensor]:
+    left: list[torch.Tensor], right: list[torch.Tensor]
+) -> list[torch.Tensor]:
     """
     Direct sum algorithm implementation to sum two tensor trains (MPS/MPO).
     It assumes the left and right bond are along the dimension 0 and -1 of each tensor.
@@ -48,8 +52,8 @@ def add_factors(
 
 
 def scale_factors(
-    factors: list[torch.tensor], scalar: complex, *, which: int
-) -> list[torch.tensor]:
+    factors: list[torch.Tensor], scalar: complex, *, which: int
+) -> list[torch.Tensor]:
     """
     Returns a new list of factors where the tensor at the given index is scaled by `scalar`.
     """
@@ -57,10 +61,10 @@ def scale_factors(
 
 
 def zip_right_step(
-    slider: torch.tensor,
-    top: torch.tensor,
-    bottom: torch.tensor,
-) -> torch.tensor:
+    slider: torch.Tensor,
+    top: torch.Tensor,
+    bottom: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Returns a new `MPS/O` factor of the result of the multiplication MPO @ MPS/O,
     and the updated slider, performing a single step of the
@@ -113,11 +117,10 @@ def zip_right_step(
 
 
 def zip_right(
-    top_factors: list[torch.tensor],
-    bottom_factors: list[torch.tensor],
-    max_error: float = 1e-5,
-    max_rank: int = 1024,
-) -> list[torch.tensor]:
+    top_factors: list[torch.Tensor],
+    bottom_factors: list[torch.Tensor],
+    config: Optional[MPSConfig] = None,
+) -> list[torch.Tensor]:
     """
     Returns a new matrix product, resulting from applying `top` to `bottom`.
     The resulting factors are:
@@ -136,6 +139,8 @@ def zip_right(
         A final truncation sweep, from right to left,
         moves back the orthogonal center to the first element.
     """
+    config = config if config is not None else MPSConfig()
+
     if len(top_factors) != len(bottom_factors):
         raise ValueError("Cannot multiply two matrix products of different lengths.")
 
@@ -146,6 +151,6 @@ def zip_right(
         new_factors.append(res)
     new_factors[-1] @= slider[:, :, 0]
 
-    truncate_impl(new_factors, max_error=max_error, max_rank=max_rank)
+    truncate_impl(new_factors, config=config)
 
     return new_factors

emu_mps/constants.py

@@ -0,0 +1,4 @@
+import torch
+
+
+DEVICE_COUNT = torch.cuda.device_count()

emu_mps/hamiltonian.py

@@ -285,6 +285,7 @@ def _get_interactions_to_keep(interaction_matrix: torch.Tensor) -> list[torch.Te
     returns a list of bool valued tensors,
     indicating which interaction terms to keep for each bond in the MPO
     """
+    interaction_matrix = interaction_matrix.clone()
     nqubits = interaction_matrix.size(dim=1)
     middle = nqubits // 2
     interaction_matrix += torch.eye(
@@ -359,10 +360,9 @@ def make_H(
 
     nqubits = interaction_matrix.size(dim=1)
     middle = nqubits // 2
-
     interactions_to_keep = _get_interactions_to_keep(interaction_matrix)
 
-    cores = [_first_factor(interactions_to_keep[0].item())]
+    cores = [_first_factor(interactions_to_keep[0].item() != 0.0)]
 
     if nqubits > 2:
         for i in range(1, middle):
@@ -394,7 +394,7 @@ def make_H(
                 )
             )
     if nqubits == 2:
-        scale = interaction_matrix[0, 1]
+        scale = interaction_matrix[0, 1].item()
     elif interactions_to_keep[-1][0]:
         scale = 1.0
     else:

emu_mps/mpo.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import itertools
-from typing import Any, List, cast, Iterable, Optional
+from typing import Any, List, Iterable, Optional
 
 import torch
 
@@ -8,7 +8,8 @@ from emu_mps.algebra import add_factors, scale_factors, zip_right
 from emu_base.base_classes.operator import FullOp, QuditOp
 from emu_base import Operator, State
 from emu_mps.mps import MPS
-from emu_mps.utils import new_left_bath, assign_devices, DEVICE_COUNT
+from emu_mps.utils import new_left_bath, assign_devices
+from emu_mps.constants import DEVICE_COUNT
 
 
 def _validate_operator_targets(operations: FullOp, nqubits: int) -> None:
@@ -80,8 +81,7 @@ class MPO(Operator):
         factors = zip_right(
             self.factors,
             other.factors,
-            max_error=other.precision,
-            max_rank=other.max_bond_dim,
+            config=other.config,
         )
         return MPS(factors, orthogonality_center=0)
 
@@ -175,14 +175,15 @@ class MPO(Operator):
         Returns:
             the operator in MPO form.
         """
+        operators_with_tensors: dict[str, torch.Tensor | QuditOp] = dict(operators)
 
         _validate_operator_targets(operations, nqubits)
 
         basis = set(basis)
         if basis == {"r", "g"}:
-            # operators will now contain the basis for single qubit ops, and potentially
-            # user defined strings in terms of these
-            operators |= {
+            # operators_with_tensors will now contain the basis for single qubit ops,
+            # and potentially user defined strings in terms of these
+            operators_with_tensors |= {
                 "gg": torch.tensor(
                     [[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128
                 ).reshape(1, 2, 2, 1),
@@ -197,9 +198,9 @@ class MPO(Operator):
                 ).reshape(1, 2, 2, 1),
             }
         elif basis == {"0", "1"}:
-            # operators will now contain the basis for single qubit ops, and potentially
-            # user defined strings in terms of these
-            operators |= {
+            # operators_with_tensors will now contain the basis for single qubit ops,
+            # and potentially user defined strings in terms of these
+            operators_with_tensors |= {
                 "00": torch.tensor(
                     [[1.0, 0.0], [0.0, 0.0]], dtype=torch.complex128
                 ).reshape(1, 2, 2, 1),
@@ -218,26 +219,28 @@ class MPO(Operator):
 
         mpos = []
         for coeff, tensorop in operations:
-            # this function will recurse through the operators, and replace any definitions
+            # this function will recurse through the operators_with_tensors,
+            # and replace any definitions
             # in terms of strings by the computed tensor
             def replace_operator_string(op: QuditOp | torch.Tensor) -> torch.Tensor:
-                if isinstance(op, dict):
-                    for opstr, coeff in op.items():
-                        tensor = replace_operator_string(operators[opstr])
-                        operators[opstr] = tensor
-                        op[opstr] = tensor * coeff
-                    op = sum(cast(list[torch.Tensor], op.values()))
-                return op
+                if isinstance(op, torch.Tensor):
+                    return op
+
+                result = torch.zeros(1, 2, 2, 1, dtype=torch.complex128)
+                for opstr, coeff in op.items():
+                    tensor = replace_operator_string(operators_with_tensors[opstr])
+                    operators_with_tensors[opstr] = tensor
+                    result += tensor * coeff
+                return result
 
             factors = [
                 torch.eye(2, 2, dtype=torch.complex128).reshape(1, 2, 2, 1)
             ] * nqubits
 
-            for i, op in enumerate(tensorop):
-                tensorop[i] = (replace_operator_string(op[0]), op[1])
-
             for op in tensorop:
-                for i in op[1]:
-                    factors[i] = op[0]
+                factor = replace_operator_string(op[0])
+                for target_qubit in op[1]:
+                    factors[target_qubit] = factor
+
             mpos.append(coeff * MPO(factors, **kwargs))
         return sum(mpos[1:], start=mpos[0])

emu_mps/mps.py

@@ -7,15 +7,16 @@ from typing import Any, List, Optional, Iterable
 import torch
 
 from emu_base import State
+from emu_mps import MPSConfig
 from emu_mps.algebra import add_factors, scale_factors
 from emu_mps.utils import (
-    DEVICE_COUNT,
     apply_measurement_errors,
     assign_devices,
     truncate_impl,
     tensor_trace,
     n_operator,
 )
+from emu_mps.constants import DEVICE_COUNT
 
 
 class MPS(State):
@@ -27,17 +28,13 @@ class MPS(State):
     Only qubits are supported.
     """
 
-    DEFAULT_MAX_BOND_DIM: int = 1024
-    DEFAULT_PRECISION: float = 1e-5
-
     def __init__(
         self,
         factors: List[torch.Tensor],
         /,
         *,
         orthogonality_center: Optional[int] = None,
-        precision: float = DEFAULT_PRECISION,
-        max_bond_dim: int = DEFAULT_MAX_BOND_DIM,
+        config: Optional[MPSConfig] = None,
         num_gpus_to_use: Optional[int] = DEVICE_COUNT,
     ):
         """
@@ -53,14 +50,11 @@ class MPS(State):
                 of the data to this constructor, or some shared objects.
             orthogonality_center: the orthogonality center of the MPS, or None (in which case
                 it will be orthogonalized when needed)
-            precision: the precision with which to truncate here or in tdvp
-            max_bond_dim: the maximum bond dimension to allow
+            config: the emu-mps config object passed to the run method
             num_gpus_to_use: distribute the factors over this many GPUs
                 0=all factors to cpu, None=keep the existing device assignment.
         """
-        self.precision = precision
-        self.max_bond_dim = max_bond_dim
-
+        self.config = config if config is not None else MPSConfig()
         assert all(
             factors[i - 1].shape[2] == factors[i].shape[0] for i in range(1, len(factors))
         ), "The dimensions of consecutive tensors should match"
@@ -84,8 +78,7 @@ class MPS(State):
     def make(
         cls,
         num_sites: int,
-        precision: float = DEFAULT_PRECISION,
-        max_bond_dim: int = DEFAULT_MAX_BOND_DIM,
+        config: Optional[MPSConfig] = None,
         num_gpus_to_use: int = DEVICE_COUNT,
     ) -> MPS:
         """
@@ -93,11 +86,12 @@ class MPS(State):
 
         Args:
             num_sites: the number of qubits
-            precision: the precision with which to truncate here or in tdvp
-            max_bond_dim: the maximum bond dimension to allow
+            config: the MPSConfig
             num_gpus_to_use: distribute the factors over this many GPUs
                 0=all factors to cpu
         """
+        config = config if config is not None else MPSConfig()
+
         if num_sites <= 1:
             raise ValueError("For 1 qubit states, do state vector")
 
@@ -106,8 +100,7 @@ class MPS(State):
                 torch.tensor([[[1.0], [0.0]]], dtype=torch.complex128)
                 for _ in range(num_sites)
             ],
-            precision=precision,
-            max_bond_dim=max_bond_dim,
+            config=config,
             num_gpus_to_use=num_gpus_to_use,
             orthogonality_center=0,  # Arbitrary: every qubit is an orthogonality center.
         )
@@ -165,15 +158,11 @@ class MPS(State):
         """
         SVD based truncation of the state. Puts the orthogonality center at the first qubit.
         Calls orthogonalize on the last qubit, and then sweeps a series of SVDs right-left.
-        Uses self.precision and self.max_bond_dim for determining accuracy.
+        Uses self.config for determining accuracy.
         An in-place operation.
         """
         self.orthogonalize(self.num_sites - 1)
-        truncate_impl(
-            self.factors,
-            max_error=self.precision,
-            max_rank=self.max_bond_dim,
-        )
+        truncate_impl(self.factors, config=self.config)
         self.orthogonality_center = 0
 
     def get_max_bond_dim(self) -> int:
@@ -236,7 +225,7 @@ class MPS(State):
         num_qubits = len(self.factors)
 
         bitstring = ""
-        acc_mps_j: torch.tensor = self.factors[0]
+        acc_mps_j: torch.Tensor = self.factors[0]
 
         for qubit in range(num_qubits):
             # comp_basis is a projector: 0 is for ket |0> and 1 for ket |1>
@@ -303,7 +292,7 @@ class MPS(State):
         """
         Returns the sum of two MPSs, computed with a direct algorithm.
         The resulting MPS is orthogonalized on the first site and truncated
-        up to `self.precision`.
+        up to `self.config.precision`.
 
         Args:
             other: the other state
@@ -315,8 +304,7 @@ class MPS(State):
         new_tt = add_factors(self.factors, other.factors)
         result = MPS(
             new_tt,
-            precision=self.precision,
-            max_bond_dim=self.max_bond_dim,
+            config=self.config,
             num_gpus_to_use=None,
             orthogonality_center=None,  # Orthogonality is lost.
         )
@@ -341,8 +329,7 @@ class MPS(State):
         factors = scale_factors(self.factors, scalar, which=which)
         return MPS(
             factors,
-            precision=self.precision,
-            max_bond_dim=self.max_bond_dim,
+            config=self.config,
             num_gpus_to_use=None,
             orthogonality_center=self.orthogonality_center,
         )

emu_mps/mps_backend.py

@@ -1,8 +1,12 @@
-from pulser import Sequence
-
 from emu_base import Backend, BackendConfig, Results
 from emu_mps.mps_config import MPSConfig
-from emu_mps.mps_backend_impl import create_impl
+from emu_mps.mps_backend_impl import create_impl, MPSBackendImpl
+from pulser import Sequence
+import pickle
+import os
+import time
+import logging
+import pathlib
 
 
 class MPSBackend(Backend):
@@ -11,6 +15,32 @@ class MPSBackend(Backend):
     aka tensor trains.
     """
 
+    def resume(self, autosave_file: str | pathlib.Path) -> Results:
+        """
+        Resume simulation from autosave file.
+        Only resume simulations from data you trust!
+        Unpickling of untrusted data is not safe.
+        """
+        if isinstance(autosave_file, str):
+            autosave_file = pathlib.Path(autosave_file)
+
+        if not autosave_file.is_file():
+            raise ValueError(f"Not a file: {autosave_file}")
+
+        with open(autosave_file, "rb") as f:
+            impl: MPSBackendImpl = pickle.load(f)
+
+        impl.autosave_file = autosave_file
+        impl.last_save_time = time.time()
+        impl.config.init_logging()  # FIXME: might be best to take logger object out of config.
+
+        logging.getLogger("global_logger").warning(
+            f"Resuming simulation from file {autosave_file}\n"
+            f"Saving simulation state every {impl.config.autosave_dt} seconds"
+        )
+
+        return self._run(impl)
+
     def run(self, sequence: Sequence, mps_config: BackendConfig) -> Results:
         """
         Emulates the given sequence.
@@ -29,7 +59,14 @@ class MPSBackend(Backend):
         impl = create_impl(sequence, mps_config)
         impl.init()  # This is separate from the constructor for testing purposes.
 
+        return self._run(impl)
+
+    @staticmethod
+    def _run(impl: MPSBackendImpl) -> Results:
         while not impl.is_finished():
             impl.progress()
 
+        if impl.autosave_file.is_file():
+            os.remove(impl.autosave_file)
+
         return impl.results

emu_mps/mps_backend_impl.py

@@ -1,7 +1,11 @@
 import math
+import pathlib
 import random
+import uuid
 from resource import RUSAGE_SELF, getrusage
 from typing import Optional
+import pickle
+import os
 
 import torch
 import time
@@ -9,6 +13,7 @@ from pulser import Sequence
 
 from emu_base import Results, State, PulserData
 from emu_base.math.brents_root_finding import BrentsRootFinder
+from emu_mps.constants import DEVICE_COUNT
 from emu_mps.hamiltonian import make_H, update_H
 from emu_mps.mpo import MPO
 from emu_mps.mps import MPS
@@ -17,7 +22,6 @@ from emu_mps.noise import compute_noise_from_lindbladians, pick_well_prepared_qu
 from emu_mps.tdvp import (
     evolve_single,
     evolve_pair,
-    EvolveConfig,
     new_right_bath,
     right_baths,
 )
@@ -52,6 +56,7 @@ class MPSBackendImpl:
     def __init__(self, mps_config: MPSConfig, pulser_data: PulserData):
         self.config = mps_config
         self.target_time = float(self.config.dt)
+        self.pulser_data = pulser_data
         self.qubit_count = pulser_data.qubit_count
         assert self.qubit_count >= 2
         self.omega = pulser_data.omega
@@ -71,15 +76,23 @@ class MPSBackendImpl:
         self.tdvp_index = 0
         self.timestep_index = 0
         self.results = Results()
-
-        self.evolve_config = EvolveConfig(
-            exp_tolerance=self.config.precision * self.config.extra_krylov_tolerance,
-            norm_tolerance=self.config.precision * self.config.extra_krylov_tolerance,
-            max_krylov_dim=self.config.max_krylov_dim,
-            is_hermitian=not self.has_lindblad_noise,
-            max_error=self.config.precision,
-            max_rank=self.config.max_bond_dim,
+        self.autosave_file = self._get_autosave_filepath(self.config.autosave_prefix)
+        self.config.logger.warning(
+            f"""Will save simulation state to file "{self.autosave_file.name}"
+            every {self.config.autosave_dt} seconds.\n"""
+            f"""To resume: `MPSBackend().resume("{self.autosave_file}")`"""
         )
+        self.last_save_time = time.time()
+
+        if self.config.num_gpus_to_use > DEVICE_COUNT:
+            self.config.logger.warning(
+                f"Requested to use {self.config.num_gpus_to_use} GPU(s) "
+                f"but only {DEVICE_COUNT if DEVICE_COUNT > 0 else 'cpu'} available"
+            )
+
+    @staticmethod
+    def _get_autosave_filepath(autosave_prefix: str) -> pathlib.Path:
+        return pathlib.Path(os.getcwd()) / (autosave_prefix + str(uuid.uuid1()) + ".dat")
 
     def init_dark_qubits(self) -> None:
         has_state_preparation_error: bool = (
@@ -112,8 +125,7 @@ class MPSBackendImpl:
         if initial_state is None:
             self.state = MPS.make(
                 self.qubit_count,
-                precision=self.config.precision,
-                max_bond_dim=self.config.max_bond_dim,
+                config=self.config,
                 num_gpus_to_use=self.config.num_gpus_to_use,
             )
             return
@@ -128,8 +140,7 @@ class MPSBackendImpl:
         initial_state = MPS(
             # Deep copy of every tensor of the initial state.
             [f.clone().detach() for f in initial_state.factors],
-            precision=self.config.precision,
-            max_bond_dim=self.config.max_bond_dim,
+            config=self.config,
             num_gpus_to_use=self.config.num_gpus_to_use,
         )
         initial_state.truncate()
@@ -198,7 +209,8 @@ class MPSBackendImpl:
                 ham_factor=self.hamiltonian.factors[index],
                 baths=baths,
                 dt=dt,
-                config=self.evolve_config,
+                config=self.config,
+                is_hermitian=not self.has_lindblad_noise,
             )
         else:
             assert orth_center_right is not None
@@ -213,8 +225,9 @@ class MPSBackendImpl:
                 ham_factors=self.hamiltonian.factors[l : r + 1],
                 baths=baths,
                 dt=dt,
-                config=self.evolve_config,
+                config=self.config,
                 orth_center_right=orth_center_right,
+                is_hermitian=not self.has_lindblad_noise,
             )
 
             self.state.orthogonality_center = r if orth_center_right else l
@@ -292,7 +305,7 @@ class MPSBackendImpl:
             )
             if not self.has_lindblad_noise:
                 # Free memory because it won't be used anymore
-                self.right_baths[-2] = None
+                self.right_baths[-2] = torch.zeros(0)
 
             self._evolve(self.tdvp_index, dt=-delta_time / 2)
             self.left_baths.pop()
@@ -312,7 +325,7 @@ class MPSBackendImpl:
         else:
             raise Exception("Didn't expect this")
 
-        # TODO: checkpoint/autosave here
+        self.save_simulation()
 
     def tdvp_complete(self) -> None:
         self.current_time = self.target_time
@@ -343,14 +356,40 @@ class MPSBackendImpl:
         self.log_step_statistics(duration=time.time() - self.time)
         self.time = time.time()
 
+    def save_simulation(self) -> None:
+        if self.last_save_time > time.time() - self.config.autosave_dt:
+            return
+
+        basename = self.autosave_file
+        with open(basename.with_suffix(".new"), "wb") as file_handle:
+            pickle.dump(self, file_handle)
+
+        if basename.is_file():
+            os.rename(basename, basename.with_suffix(".bak"))
+
+        os.rename(basename.with_suffix(".new"), basename)
+        autosave_filesize = os.path.getsize(self.autosave_file) / 1e6
+
+        if basename.with_suffix(".bak").is_file():
+            os.remove(basename.with_suffix(".bak"))
+
+        self.last_save_time = time.time()
+
+        self.config.logger.debug(
+            f"Saved simulation state in file {self.autosave_file} ({autosave_filesize}MB)"
+        )
+
     def fill_results(self) -> None:
         normalized_state = 1 / self.state.norm() * self.state
 
+        current_time_int: int = round(self.current_time)
+        assert abs(self.current_time - current_time_int) < 1e-10
+
         if self.well_prepared_qubits_filter is None:
             for callback in self.config.callbacks:
                 callback(
                     self.config,
-                    self.current_time,
+                    current_time_int,
                     normalized_state,
                     self.hamiltonian,
                     self.results,
@@ -359,7 +398,7 @@ class MPSBackendImpl:
 
         full_mpo, full_state = None, None
         for callback in self.config.callbacks:
-            if self.current_time not in callback.evaluation_times:
+            if current_time_int not in callback.evaluation_times:
                 continue
 
             if full_mpo is None or full_state is None:
@@ -380,7 +419,7 @@ class MPSBackendImpl:
                     ),
                 )
 
-            callback(self.config, self.current_time, full_state, full_mpo, self.results)
+            callback(self.config, current_time_int, full_state, full_mpo, self.results)
 
     def log_step_statistics(self, *, duration: float) -> None:
         if self.state.factors[0].is_cuda:
@@ -424,13 +463,12 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
     """
 
     jump_threshold: float
-    aggregated_lindblad_ops: Optional[torch.Tensor]
+    aggregated_lindblad_ops: torch.Tensor
     norm_gap_before_jump: float
     root_finder: Optional[BrentsRootFinder]
 
     def __init__(self, config: MPSConfig, pulser_data: PulserData):
         super().__init__(config, pulser_data)
-        self.aggregated_lindblad_ops = None
         self.lindblad_ops = pulser_data.lindblad_ops
         self.root_finder = None
 
@@ -491,7 +529,7 @@ class NoisyMPSBackendImpl(MPSBackendImpl):
                 for qubit in range(self.state.num_sites)
                 for op in self.lindblad_ops
             ],
-            weights=jump_operator_weights.reshape(-1),
+            weights=jump_operator_weights.reshape(-1).tolist(),
         )[0]
 
         self.state.apply(jumped_qubit_index, jump_operator)

emu_mps/mps_config.py

@@ -1,7 +1,7 @@
 from typing import Any
 
 from emu_base import BackendConfig, State
-from emu_mps.utils import DEVICE_COUNT
+from emu_mps.constants import DEVICE_COUNT
 
 
 class MPSConfig(BackendConfig):
@@ -23,6 +23,10 @@ class MPSConfig(BackendConfig):
         num_gpus_to_use: during the simulation, distribute the state over this many GPUs
             0=all factors to cpu. As shown in the benchmarks, using multiple GPUs might
             alleviate memory pressure per GPU, but the runtime should be similar.
+        autosave_prefix: filename prefix for autosaving simulation state to file
+        autosave_dt: minimum time interval in seconds between two autosaves
+            Saving the simulation state is only possible at specific times,
+            therefore this interval is only a lower bound.
         kwargs: arguments that are passed to the base class
 
     Examples:
@@ -43,6 +47,8 @@ class MPSConfig(BackendConfig):
         max_krylov_dim: int = 100,
         extra_krylov_tolerance: float = 1e-3,
         num_gpus_to_use: int = DEVICE_COUNT,
+        autosave_prefix: str = "emu_mps_save_",
+        autosave_dt: int = 600,  # 10 minutes
         **kwargs: Any,
     ):
         super().__init__(**kwargs)
@@ -62,3 +68,11 @@ class MPSConfig(BackendConfig):
                 and self.noise_model.amp_sigma != 0.0
             ):
                 raise NotImplementedError("Unsupported noise type: amp_sigma")
+
+        self.autosave_prefix = autosave_prefix
+        self.autosave_dt = autosave_dt
+
+        MIN_AUTOSAVE_DT = 10
+        assert (
+            self.autosave_dt > MIN_AUTOSAVE_DT
+        ), f"autosave_dt must be larger than {MIN_AUTOSAVE_DT} seconds"