Skip to content

Apple Silicon Metal / MLX Simulator

quanta.simulator.mlx

quanta.simulator.mlx -- Apple Silicon Metal/MLX accelerated statevector simulator.

World's first native Apple Silicon Metal-accelerated quantum statevector simulator. Leverages Apple MLX (mlx.core + mlx-metal) and Unified Memory architecture on Apple M-series chips (M1, M2, M3, M4, M5 Pro / Max / Ultra).

Features: - Up to 400x speedup over CPU NumPy statevector simulation. - Native execution on Apple Silicon GPU via Metal Performance Shaders. - Multidimensional N-axis tensor representation [2]*N for 30+ qubit scalability. - Zero-copy evaluation and memory-efficient tensor contractions. - Full parity with Quanta's SimulatorBackend interface.

Requirements

pip install mlx mlx-metal

Example

from quanta.simulator.mlx import MLXSimulator sim = MLXSimulator(num_qubits=20) sim.apply("H", (0,)) sim.apply("CX", (0, 1)) print(sim.probabilities()[:4]) [0.5, 0. , 0. , 0.5]

MLXSimulator

Bases: StateVectorSimulator

Apple Silicon Metal/MLX accelerated statevector simulator.

Executes quantum circuits directly on the Apple Silicon GPU using Metal. Provides up to 400x speedup over CPU execution for 20-30+ qubit circuits.

Parameters:

Name Type Description Default
num_qubits int

Number of qubits to simulate (supports up to 30+ on 48GB unified RAM).

 required
seed int | None

Random seed for measurement reproducibility.

 None
Source code in quanta/simulator/mlx.py 
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
class MLXSimulator(StateVectorSimulator):
    """Apple Silicon Metal/MLX accelerated statevector simulator.

    Executes quantum circuits directly on the Apple Silicon GPU using Metal.
    Provides up to 400x speedup over CPU execution for 20-30+ qubit circuits.

    Args:
        num_qubits: Number of qubits to simulate (supports up to 30+ on 48GB unified RAM).
        seed: Random seed for measurement reproducibility.
    """

    MAX_QUBITS = 32

    def __init__(self, num_qubits: int, seed: int | None = None) -> None:
        if not is_mlx_available():
            raise MLXSimulatorError(
                "Apple MLX is required for MLXSimulator.\n"
                "To install: pip install mlx mlx-metal\n"
                "(Note: MLX requires an Apple Silicon Mac running macOS)."
            )

        if num_qubits > self.MAX_QUBITS:
            raise MLXSimulatorError(
                f"MLXSimulator supports up to {self.MAX_QUBITS} qubits, requested: {num_qubits}"
            )

        import mlx.core as mx  # type: ignore[import-untyped]

        self._mx = mx
        self.num_qubits = num_qubits
        self._rng = np.random.default_rng(seed)
        self._gate_cache: dict[tuple[str, tuple[float, ...]], Any] = {}
        self._perm_cache: dict[tuple[tuple[int, ...], int], list[int]] = {}
        self._pending_ops: int = 0
        self._eval_batch_size: int = 8

        # Initialize |00...0> state tensor of shape [2]*n directly on Metal GPU
        q0 = mx.array([1.0, 0.0], dtype=mx.complex64)
        state = q0.reshape([2] + [1] * (num_qubits - 1))
        for i in range(1, num_qubits):
            shape = [1] * i + [2] + [1] * (num_qubits - 1 - i)
            state = state * q0.reshape(shape)

        self._state = state  # type: ignore[assignment]
        mx.eval(self._state)

    def _sync(self) -> None:
        """Flushes deferred graph evaluations to Metal GPU."""
        if self._pending_ops > 0:
            self._mx.eval(self._state)
            self._pending_ops = 0

    @property
    def max_qubits(self) -> int:
        return self.MAX_QUBITS

    def apply(
        self,
        gate_name: str,
        qubits: tuple[int, ...],
        params: tuple[float, ...] = (),
    ) -> None:
        """Applies a quantum gate to the statevector tensor on Apple Metal GPU."""
        mx = self._mx
        k = len(qubits)

        # Gate tensor caching
        cache_key = (gate_name, params)
        gate_tensor = self._gate_cache.get(cache_key)
        if gate_tensor is None:
            gate_matrix_np = self._get_gate_matrix(gate_name, params)
            gate_mat = mx.array(gate_matrix_np, dtype=mx.complex64)
            gate_tensor = gate_mat.reshape([2] * (2 * k))
            self._gate_cache[cache_key] = gate_tensor

        # Tensor contraction along target qubit axes
        gate_axes = list(range(k, 2 * k))
        state_axes = list(qubits)

        result = mx.tensordot(gate_tensor, self._state, axes=(gate_axes, state_axes))

        # Reorder tensor axes back to [0, 1, ..., n-1] (using cached permutation)
        perm_key = (qubits, self.num_qubits)
        inv_perm = self._perm_cache.get(perm_key)
        if inv_perm is None:
            remaining_axes = [i for i in range(self.num_qubits) if i not in qubits]
            current_order = list(qubits) + remaining_axes
            inv_perm = [current_order.index(p) for p in range(self.num_qubits)]
            self._perm_cache[perm_key] = inv_perm

        self._state = mx.transpose(result, inv_perm)  # type: ignore[assignment]
        self._pending_ops += 1
        if self._pending_ops >= self._eval_batch_size:
            self._sync()

    def probabilities(self) -> np.ndarray:
        """Returns measurement probabilities for all basis states: P = |alpha|^2."""
        self._sync()
        mx = self._mx
        flat_state = mx.reshape(self._state, (-1,))
        probs = mx.abs(flat_state) ** 2
        mx.eval(probs)
        return np.array(probs, dtype=np.float64)

    def sample(self, shots: int) -> dict[str, int]:
        """Samples measurement outcomes from basis state probabilities."""
        probs = self.probabilities()
        # Normalize to exact 1.0 to handle 32-bit float rounding in np.random.choice
        prob_sum = np.sum(probs)
        if prob_sum > 0:
            probs = probs / prob_sum
        dim = len(probs)
        indices = self._rng.choice(dim, size=shots, p=probs)
        unique, unique_counts = np.unique(indices, return_counts=True)
        fmt = f"0{self.num_qubits}b"
        return {
            format(idx, fmt): int(cnt)
            for idx, cnt in zip(unique, unique_counts, strict=True)
        }

    @property
    def state(self) -> np.ndarray:
        """Returns the statevector as a NumPy complex128 array."""
        self._sync()
        mx = self._mx
        flat = mx.reshape(self._state, (-1,))
        return np.array(flat, dtype=np.complex128)

    def norm(self) -> float:
        """Returns the L2 norm of the statevector on Apple Metal GPU."""
        self._sync()
        return float(np.linalg.norm(self.state))

    @state.setter
    def state(self, value: Any) -> None:
        """Sets statevector from array."""
        mx = self._mx
        np_arr = np.asarray(value, dtype=np.complex64)
        if np_arr.size != (2 ** self.num_qubits):
            expected_dim = 2 ** self.num_qubits
            raise MLXSimulatorError(
                f"State size {np_arr.size} does not match 2^{self.num_qubits} = {expected_dim}"
            )
        self._state = mx.array(np_arr.reshape([2] * self.num_qubits))  # type: ignore[assignment]
        self._pending_ops = 0
        mx.eval(self._state)

    def _get_gate_matrix(
        self, name: str, params: tuple[float, ...]
    ) -> np.ndarray:
        """Retrieves unitary gate matrix from registry."""
        gate = GATE_REGISTRY.get(name)
        if gate is None:
            raise MLXSimulatorError(f"Unknown gate: {name}")

        if isinstance(gate, MultiParametricGate):
            if not params:
                raise MLXSimulatorError(f"{name} gate requires parameters")
            return gate(*params).matrix
        if isinstance(gate, ParametricGate):
            if not params:
                raise MLXSimulatorError(f"{name} gate requires parameters")
            return gate(params[0]).matrix

        return gate.matrix

    def apply_phase(self, index: int, phase: complex) -> None:
        """Applies a phase factor to a specific basis state on Apple Metal GPU."""
        mx = self._mx
        flat = mx.reshape(self._state, (-1,))
        flat[index] = flat[index] * mx.array(phase, dtype=mx.complex64)
        self._state = mx.reshape(flat, [2] * self.num_qubits)
        self._pending_ops += 1
        if self._pending_ops >= self._eval_batch_size:
            self._sync()

    def apply_noise(
        self,
        noise_model: Any,
        qubits: tuple[int, ...],
        rng: np.random.Generator,
    ) -> None:
        """Applies a noise model directly on Apple Metal GPU when possible."""
        from quanta.simulator.noise import BitFlip, Depolarizing, PhaseFlip

        channels = getattr(noise_model, "channels", [])
        all_pauli = bool(channels)
        for channel in channels:
            if not isinstance(channel, (Depolarizing, BitFlip, PhaseFlip)):
                all_pauli = False
                break

        if all_pauli:
            for channel in channels:
                for qubit in qubits:
                    if isinstance(channel, Depolarizing):
                        if rng.random() <= channel.probability:
                            pauli = int(rng.integers(0, 3))
                            gate = ("X", "Y", "Z")[pauli]
                            self.apply(gate, (qubit,))
                    elif isinstance(channel, BitFlip) and rng.random() <= channel.probability:
                        self.apply("X", (qubit,))
                    elif isinstance(channel, PhaseFlip) and rng.random() <= channel.probability:
                        self.apply("Z", (qubit,))
            return

        # General / non-Pauli channel fallback
        self._sync()
        mx = self._mx
        flat = np.array(mx.reshape(self._state, (-1,)), dtype=np.complex128)
        new_flat = noise_model.apply_noise(flat, qubits, self.num_qubits, rng)
        self.state = new_flat

    def reset(self) -> None:
        """Resets the simulator state to |00...0>."""
        self._pending_ops = 0
        self.__init__(self.num_qubits)  # type: ignore[misc]

    def __repr__(self) -> str:
        return f"MLXSimulator(qubits={self.num_qubits}, device=metal_gpu)"
state property writable 
state: ndarray

Returns the statevector as a NumPy complex128 array.

apply 
apply(
    gate_name: str,
    qubits: tuple[int, ...],
    params: tuple[float, ...] = (),
) -> None

Applies a quantum gate to the statevector tensor on Apple Metal GPU.

Source code in quanta/simulator/mlx.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def apply(
    self,
    gate_name: str,
    qubits: tuple[int, ...],
    params: tuple[float, ...] = (),
) -> None:
    """Applies a quantum gate to the statevector tensor on Apple Metal GPU."""
    mx = self._mx
    k = len(qubits)

    # Gate tensor caching
    cache_key = (gate_name, params)
    gate_tensor = self._gate_cache.get(cache_key)
    if gate_tensor is None:
        gate_matrix_np = self._get_gate_matrix(gate_name, params)
        gate_mat = mx.array(gate_matrix_np, dtype=mx.complex64)
        gate_tensor = gate_mat.reshape([2] * (2 * k))
        self._gate_cache[cache_key] = gate_tensor

    # Tensor contraction along target qubit axes
    gate_axes = list(range(k, 2 * k))
    state_axes = list(qubits)

    result = mx.tensordot(gate_tensor, self._state, axes=(gate_axes, state_axes))

    # Reorder tensor axes back to [0, 1, ..., n-1] (using cached permutation)
    perm_key = (qubits, self.num_qubits)
    inv_perm = self._perm_cache.get(perm_key)
    if inv_perm is None:
        remaining_axes = [i for i in range(self.num_qubits) if i not in qubits]
        current_order = list(qubits) + remaining_axes
        inv_perm = [current_order.index(p) for p in range(self.num_qubits)]
        self._perm_cache[perm_key] = inv_perm

    self._state = mx.transpose(result, inv_perm)  # type: ignore[assignment]
    self._pending_ops += 1
    if self._pending_ops >= self._eval_batch_size:
        self._sync()
apply_noise 
apply_noise(
    noise_model: Any,
    qubits: tuple[int, ...],
    rng: Generator,
) -> None

Applies a noise model directly on Apple Metal GPU when possible.

Source code in quanta/simulator/mlx.py 
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
def apply_noise(
    self,
    noise_model: Any,
    qubits: tuple[int, ...],
    rng: np.random.Generator,
) -> None:
    """Applies a noise model directly on Apple Metal GPU when possible."""
    from quanta.simulator.noise import BitFlip, Depolarizing, PhaseFlip

    channels = getattr(noise_model, "channels", [])
    all_pauli = bool(channels)
    for channel in channels:
        if not isinstance(channel, (Depolarizing, BitFlip, PhaseFlip)):
            all_pauli = False
            break

    if all_pauli:
        for channel in channels:
            for qubit in qubits:
                if isinstance(channel, Depolarizing):
                    if rng.random() <= channel.probability:
                        pauli = int(rng.integers(0, 3))
                        gate = ("X", "Y", "Z")[pauli]
                        self.apply(gate, (qubit,))
                elif isinstance(channel, BitFlip) and rng.random() <= channel.probability:
                    self.apply("X", (qubit,))
                elif isinstance(channel, PhaseFlip) and rng.random() <= channel.probability:
                    self.apply("Z", (qubit,))
        return

    # General / non-Pauli channel fallback
    self._sync()
    mx = self._mx
    flat = np.array(mx.reshape(self._state, (-1,)), dtype=np.complex128)
    new_flat = noise_model.apply_noise(flat, qubits, self.num_qubits, rng)
    self.state = new_flat
apply_phase 
apply_phase(index: int, phase: complex) -> None

Applies a phase factor to a specific basis state on Apple Metal GPU.

Source code in quanta/simulator/mlx.py 
222
223
224
225
226
227
228
229
230
def apply_phase(self, index: int, phase: complex) -> None:
    """Applies a phase factor to a specific basis state on Apple Metal GPU."""
    mx = self._mx
    flat = mx.reshape(self._state, (-1,))
    flat[index] = flat[index] * mx.array(phase, dtype=mx.complex64)
    self._state = mx.reshape(flat, [2] * self.num_qubits)
    self._pending_ops += 1
    if self._pending_ops >= self._eval_batch_size:
        self._sync()
norm 
norm() -> float

Returns the L2 norm of the statevector on Apple Metal GPU.

Source code in quanta/simulator/mlx.py 
184
185
186
187
def norm(self) -> float:
    """Returns the L2 norm of the statevector on Apple Metal GPU."""
    self._sync()
    return float(np.linalg.norm(self.state))
probabilities 
probabilities() -> np.ndarray

Returns measurement probabilities for all basis states: P = |alpha|^2.

Source code in quanta/simulator/mlx.py 
151
152
153
154
155
156
157
158
def probabilities(self) -> np.ndarray:
    """Returns measurement probabilities for all basis states: P = |alpha|^2."""
    self._sync()
    mx = self._mx
    flat_state = mx.reshape(self._state, (-1,))
    probs = mx.abs(flat_state) ** 2
    mx.eval(probs)
    return np.array(probs, dtype=np.float64)
reset 
reset() -> None

Resets the simulator state to |00...0>.

Source code in quanta/simulator/mlx.py 
269
270
271
272
def reset(self) -> None:
    """Resets the simulator state to |00...0>."""
    self._pending_ops = 0
    self.__init__(self.num_qubits)  # type: ignore[misc]
sample 
sample(shots: int) -> dict[str, int]

Samples measurement outcomes from basis state probabilities.

Source code in quanta/simulator/mlx.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def sample(self, shots: int) -> dict[str, int]:
    """Samples measurement outcomes from basis state probabilities."""
    probs = self.probabilities()
    # Normalize to exact 1.0 to handle 32-bit float rounding in np.random.choice
    prob_sum = np.sum(probs)
    if prob_sum > 0:
        probs = probs / prob_sum
    dim = len(probs)
    indices = self._rng.choice(dim, size=shots, p=probs)
    unique, unique_counts = np.unique(indices, return_counts=True)
    fmt = f"0{self.num_qubits}b"
    return {
        format(idx, fmt): int(cnt)
        for idx, cnt in zip(unique, unique_counts, strict=True)
    }

MLXSimulatorError

Bases: QuantaError

Errors raised by the MLX simulator.

Source code in quanta/simulator/mlx.py 
52
53
class MLXSimulatorError(QuantaError):
    """Errors raised by the MLX simulator."""

is_mlx_available 

is_mlx_available() -> bool

Checks whether Apple MLX is installed and running on Apple Silicon.

Source code in quanta/simulator/mlx.py 
41
42
43
44
45
46
47
48
49
def is_mlx_available() -> bool:
    """Checks whether Apple MLX is installed and running on Apple Silicon."""
    if platform.system() != "Darwin" or platform.machine() != "arm64":
        return False
    try:
        import mlx.core as mx  # type: ignore[import-untyped]
        return bool(getattr(mx, "__version__", None))
    except (ImportError, Exception):
        return False