Run strategy="active_statevector" in numba kernels, without JAX

CHANGELOG.md

@@ -31,6 +31,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - The samplers now default to `strategy="active_statevector"`; pass
   `strategy="cat5"` for the previous stabilizer-rank default.
   `CompiledStateProbs` keeps `"cat5"`. (#203)
+- `strategy="active_statevector"` no longer uses JAX: it samples on the CPU
+  with numba kernels, also when JAX has a GPU, and trajectory sharing is on
+  by default (`trajectory_sharing=None` means True for this strategy and
+  False for the others). Amplitudes are complex128. Postselected
+  detectors of a component prune its shots with sharing off as well, and also
+  when no direct detector is postselected. (#199)
+- Without a `batch_size`, a statevector call is split into one batch per
+  numba thread whenever each batch gets at least 4096 shots, and smaller
+  calls run as one batch on the parallel kernels. An explicit `batch_size`
+  keeps samples independent of the thread count. (#199)
+- Noise sampling runs in numba kernels, so the samples for a given seed
+  differ from earlier versions; `numba>=0.61` is now a dependency. (#199)
 - Tsim now uses zxparam instead of pyzx-param. Circuits are built directly as quantum/classical ZX diagrams (no more composing with the adjoint), and reduction and stabilizer-rank decomposition run in zxparam.
 - `Circuit.get_graph()` and `Circuit.get_sampling_graph()` return `zxparam.Graph`.
 - In `Circuit.get_graph()` (and the `"zx"` diagram), every detector and observable has a classical output boundary, with ports after the qubit outputs, so the graph's tensor is the circuit's channel together with the detector and observable bits.

pyproject.toml

@@ -37,6 +37,7 @@ dependencies = [
     "ipywidgets>=8.0.0",
     "jax>=0.4.38",
     "lxml>=5.0.0",
+    "numba>=0.61",
     "numpy>=1.25.0",
     "psutil>=5.9.0",
     "stim>=1.0.0",

src/tsim/circuit.py

@@ -865,17 +865,20 @@ def compile_sampler(
         *,
         strategy: SamplingStrategy = "active_statevector",
         seed: int | None = None,
-        trajectory_sharing: bool = False,
+        trajectory_sharing: bool | None = None,
     ) -> CompiledMeasurementSampler:
         """Compile circuit into a measurement sampler.
 
         Args:
             strategy: "active_statevector" (the default) to simulate each
                 non-direct component on a dense state vector over its active
                 qubits, or a stabilizer rank decomposition strategy ("cat5",
-                "bss", "cutting").
+                "bss", "cutting"). "active_statevector" samples on the CPU, up
+                to ``numba.get_num_threads()`` batches at once, each with its
+                own memory.
             trajectory_sharing: With ``strategy="active_statevector"``, simulate shots
-                with equal histories on one lane (experimental).
+                with equal histories on one lane. Defaults to True for
+                ``"active_statevector"`` and False otherwise.
             seed: Random seed for the sampler. IMPORTANT: Currently, the sampler will
                 only produce deterministic samples for fixed batch size. If
                 deterministic samples are needed, the batch size should be set
@@ -896,7 +899,7 @@ def compile_detector_sampler(
         *,
         strategy: SamplingStrategy = "active_statevector",
         seed: int | None = None,
-        trajectory_sharing: bool = False,
+        trajectory_sharing: bool | None = None,
     ) -> CompiledDetectorSampler:
         """Compile circuit into a detector sampler.
 
@@ -909,9 +912,12 @@ def compile_detector_sampler(
             strategy: "active_statevector" (the default) to simulate each
                 non-direct component on a dense state vector over its active
                 qubits, or a stabilizer rank decomposition strategy ("cat5",
-                "bss", "cutting").
+                "bss", "cutting"). "active_statevector" samples on the CPU, up
+                to ``numba.get_num_threads()`` batches at once, each with its
+                own memory.
             trajectory_sharing: With ``strategy="active_statevector"``, simulate shots
-                with equal histories on one lane (experimental).
+                with equal histories on one lane. Defaults to True for
+                ``"active_statevector"`` and False otherwise.
             seed: Random seed for the sampler. IMPORTANT: Currently, the sampler will
                 only produce deterministic samples for fixed batch size. If
                 deterministic samples are needed, the batch size should be set

src/tsim/compile/pipeline.py

@@ -107,10 +107,10 @@ def compile_program(
 
     return CompiledProgram(
         components=tuple(compiled_components),
-        direct_f_indices=jnp.array([e[1] for e in direct_entries], dtype=jnp.int32),
-        direct_flips=jnp.array([e[2] for e in direct_entries], dtype=jnp.bool_),
-        output_order=jnp.asarray(output_order),
-        output_reindex=None if is_identity else jnp.asarray(reindex),
+        direct_f_indices=np.array([e[1] for e in direct_entries], dtype=np.int32),
+        direct_flips=np.array([e[2] for e in direct_entries], dtype=np.bool_),
+        output_order=output_order,
+        output_reindex=None if is_identity else reindex,
         num_outputs=prepared.num_outputs,
         num_detectors=prepared.num_detectors,
     )

src/tsim/core/types.py

@@ -111,9 +111,9 @@ class CompiledProgram:
     """
 
     components: tuple[CompiledComponent | StatevectorComponent, ...]
-    direct_f_indices: Array
-    direct_flips: Array
-    output_order: Array
-    output_reindex: Array | None
+    direct_f_indices: np.ndarray
+    direct_flips: np.ndarray
+    output_order: np.ndarray
+    output_reindex: np.ndarray | None
     num_outputs: int
     num_detectors: int