Skip to content

Improve sampling speed by reducing data conversion overhead #1117

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Strilanc merged 6 commits into from
Oct 9, 2026
Merged

Improve sampling speed by reducing data conversion overhead #1117

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
92e3fc0
Add `transposed` option to `stim.CompiledDetectorSampler.sample`
Strilanc Oct 3, 2026
130839c
Add unit test
Strilanc Oct 3, 2026
e764162
add baseline benchmark
Strilanc Oct 3, 2026
8c200f9
improve transposed_simd_bit_table_to_numpy_uint8 to work with bytes
Strilanc Oct 3, 2026
87df902
review
Strilanc Oct 3, 2026
d389bbf
regen
Strilanc Oct 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
140 changes: 89 additions & 51 deletions doc/python_api_reference_vDev.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6124,10 +6124,9 @@ def sample(
self,
shots: int,
*,
prepend_observables: bool = False,
append_observables: bool = False,
separate_observables: bool = False,
bit_packed: bool = False,
transposed: bool = False,
dets_out: Optional[np.ndarray] = None,
obs_out: Optional[np.ndarray] = None,
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
Expand All @@ -6142,78 +6141,117 @@ def sample(
separate_observables: Defaults to False. When set to True, the return value
is a (detection_events, observable_flips) tuple instead of a flat
detection_events array.
prepend_observables: Defaults to false. When set, observables are included
with the detectors and are placed at the start of the results.
append_observables: Defaults to false. When set, observables are included
with the detectors and are placed at the end of the results.
bit_packed: Returns a uint8 numpy array with 8 bits per byte, instead of
a bool_ numpy array with 1 bit per byte. Uses little endian packing.
transposed: Defaults to false. When set to true, outputs are transposed
(i.e. the first axis is detectors/observables and the second axis is
shots instead of vice versa).

Using transposed=True has better performance than transposed=False,
because the output more closely corresponds to stim's internal
representation. But data in this transposed form is often more difficult
to work with.
dets_out: Defaults to None. Specifies a pre-allocated numpy array to write
the detection event data into. This array must have the correct shape
and dtype.
obs_out: Defaults to None. Specifies a pre-allocated numpy array to write
the observable flip data into. This array must have the correct shape
and dtype.

Returns:
A numpy array or tuple of numpy arrays containing the samples.

if separate_observables=False and bit_packed=False:
A single numpy array.
dtype=bool_
shape=(
shots,
num_detectors + num_observables * (
append_observables + prepend_observables),
)
The bit for detection event `m` in shot `s` is at
result[s, m]
Deprecated Arguments:
(These arguments still work but are hidden and not recommended.)

prepend_observables: Prefer `separate_observables`. Defaults to false.
When set, observables are included with the detectors and are placed at
the start of the results. Using this argument is generally a bad idea,
compared to using `separate_observables=True`, because you will
inevitably just end up immediately separating the observable data
from the detector data anyways.
append_observables: Prefer `separate_observables`. Defaults to false.
When set, observables are included with the detectors and are placed at
the end of the results. Using this argument is generally a bad idea,
compared to using `separate_observables=True`, because you will
inevitably just end up immediately separating the observable data
from the detector data anyways.

if separate_observables=False and bit_packed=True:
A single numpy array.
dtype=uint8
shape=(
shots,
math.ceil((num_detectors + num_observables * (
append_observables + prepend_observables)) / 8),
)
The bit for detection event `m` in shot `s` is at
(result[s, m // 8] >> (m % 8)) & 1
Returns:
If separate_observables:
A tuple `(dets, obs)` of two numpy arrays.
else:
A single numpy array `dets`.

if separate_observables=True and bit_packed=False:
A (dets, obs) tuple.
if bit_packed=False and transposed=False:
dets.dtype=bool_
dets.shape=(shots, num_detectors)
Detection event `m` in shot `s` is dets[s, m]
if bit_packed=False and transposed=True:
dets.dtype=bool_
dets.shape=(num_detectors, shots)
Detection event `m` in shot `s` is dets[m, s]
if bit_packed=True and transposed=False:
dets.dtype=uint8
dets.shape=(shots, (num_detectors + 7) // 8)
Detection event `m` in shot `s` is `(dets[s, m // 8] >> (m % 8)) & 1`
if bit_packed=True and transposed=True:
dets.dtype=uint8
dets.shape=(num_detectors, (shots + 7) // 8)
Detection event `m` in shot `s` is `(dets[m, s // 8] >> (s % 8)) & 1`

if separate_observables and bit_packed=False and transposed=False:
obs.dtype=bool_
obs.shape=(shots, num_observables)
The bit for detection event `m` in shot `s` is at
dets[s, m]
The bit for observable `m` in shot `s` is at
obs[s, m]

if separate_observables=True and bit_packed=True:
A (dets, obs) tuple.
dets.dtype=uint8
dets.shape=(shots, math.ceil(num_detectors / 8))
Observable `m` in shot `s` is obs[s, m]
if separate_observables and bit_packed=False and transposed=True:
obs.dtype=bool_
obs.shape=(num_observables, shots)
Observable `m` in shot `s` is obs[m, s]
if separate_observables and bit_packed=True and transposed=False:
obs.dtype=uint8
obs.shape=(shots, math.ceil(num_observables / 8))
The bit for detection event `m` in shot `s` is at
(dets[s, m // 8] >> (m % 8)) & 1
The bit for observable `m` in shot `s` is at
(obs[s, m // 8] >> (m % 8)) & 1
obs.shape=(shots, (num_observables + 7) // 8)
Observable `m` in shot `s` is `(obs[s, m // 8] >> (m % 8)) & 1`
if separate_observables and bit_packed=True and transposed=True:
obs.dtype=uint8
obs.shape=(num_observables, (shots + 7) // 8)
Observable `m` in shot `s` is `(obs[m, s // 8] >> (s % 8)) & 1`

Examples:
>>> import stim
>>> c = stim.Circuit('''
... H 0
... CNOT 0 1
... R 0 1 2
... CX 0 1 2 1
... X_ERROR(1.0) 0
... M 0 1
... DETECTOR rec[-1] rec[-2]
... M 0 1 2
... DETECTOR rec[-2]
... DETECTOR rec[-1] rec[-2] rec[-3]
... OBSERVABLE_INCLUDE(0) rec[-1]
... ''')
>>> s = c.compile_detector_sampler()
>>> s.sample(shots=1)
array([[ True]])
>>> dets, obs = s.sample(shots=5, separate_observables=True)
>>> dets
array([[False, True],
[False, True],
[False, True],
[False, True],
[False, True]])
>>> obs
array([[False],
[False],
[False],
[False],
[False]])

>>> dets, obs = s.sample(
... shots=90,
... bit_packed=True,
... transposed=True,
... separate_observables=True,
... )
>>> dets
array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 3]],
dtype=uint8)
>>> obs
array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8)
"""
```

Expand Down
142 changes: 90 additions & 52 deletions doc/stim.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4763,10 +4763,9 @@ class CompiledDetectorSampler:
self,
shots: int,
*,
prepend_observables: bool = False,
append_observables: bool = False,
separate_observables: bool = False,
bit_packed: bool = False,
transposed: bool = False,
dets_out: Optional[np.ndarray] = None,
obs_out: Optional[np.ndarray] = None,
) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
Expand All @@ -4781,78 +4780,117 @@ class CompiledDetectorSampler:
separate_observables: Defaults to False. When set to True, the return value
is a (detection_events, observable_flips) tuple instead of a flat
detection_events array.
prepend_observables: Defaults to false. When set, observables are included
with the detectors and are placed at the start of the results.
append_observables: Defaults to false. When set, observables are included
with the detectors and are placed at the end of the results.
bit_packed: Returns a uint8 numpy array with 8 bits per byte, instead of
a bool_ numpy array with 1 bit per byte. Uses little endian packing.
transposed: Defaults to false. When set to true, outputs are transposed
(i.e. the first axis is detectors/observables and the second axis is
shots instead of vice versa).

Using transposed=True has better performance than transposed=False,
because the output more closely corresponds to stim's internal
representation. But data in this transposed form is often more difficult
to work with.
dets_out: Defaults to None. Specifies a pre-allocated numpy array to write
the detection event data into. This array must have the correct shape
and dtype.
obs_out: Defaults to None. Specifies a pre-allocated numpy array to write
the observable flip data into. This array must have the correct shape
and dtype.

Returns:
A numpy array or tuple of numpy arrays containing the samples.

if separate_observables=False and bit_packed=False:
A single numpy array.
dtype=bool_
shape=(
shots,
num_detectors + num_observables * (
append_observables + prepend_observables),
)
The bit for detection event `m` in shot `s` is at
result[s, m]

if separate_observables=False and bit_packed=True:
A single numpy array.
dtype=uint8
shape=(
shots,
math.ceil((num_detectors + num_observables * (
append_observables + prepend_observables)) / 8),
)
The bit for detection event `m` in shot `s` is at
(result[s, m // 8] >> (m % 8)) & 1
Deprecated Arguments:
(These arguments still work but are hidden and not recommended.)

prepend_observables: Prefer `separate_observables`. Defaults to false.
When set, observables are included with the detectors and are placed at
the start of the results. Using this argument is generally a bad idea,
compared to using `separate_observables=True`, because you will
inevitably just end up immediately separating the observable data
from the detector data anyways.
append_observables: Prefer `separate_observables`. Defaults to false.
When set, observables are included with the detectors and are placed at
the end of the results. Using this argument is generally a bad idea,
compared to using `separate_observables=True`, because you will
inevitably just end up immediately separating the observable data
from the detector data anyways.

Returns:
If separate_observables:
A tuple `(dets, obs)` of two numpy arrays.
else:
A single numpy array `dets`.

if separate_observables=True and bit_packed=False:
A (dets, obs) tuple.
if bit_packed=False and transposed=False:
dets.dtype=bool_
dets.shape=(shots, num_detectors)
Detection event `m` in shot `s` is dets[s, m]
if bit_packed=False and transposed=True:
dets.dtype=bool_
dets.shape=(num_detectors, shots)
Detection event `m` in shot `s` is dets[m, s]
if bit_packed=True and transposed=False:
dets.dtype=uint8
dets.shape=(shots, (num_detectors + 7) // 8)
Detection event `m` in shot `s` is `(dets[s, m // 8] >> (m % 8)) & 1`
if bit_packed=True and transposed=True:
dets.dtype=uint8
dets.shape=(num_detectors, (shots + 7) // 8)
Detection event `m` in shot `s` is `(dets[m, s // 8] >> (s % 8)) & 1`

if separate_observables and bit_packed=False and transposed=False:
obs.dtype=bool_
obs.shape=(shots, num_observables)
The bit for detection event `m` in shot `s` is at
dets[s, m]
The bit for observable `m` in shot `s` is at
obs[s, m]

if separate_observables=True and bit_packed=True:
A (dets, obs) tuple.
dets.dtype=uint8
dets.shape=(shots, math.ceil(num_detectors / 8))
Observable `m` in shot `s` is obs[s, m]
if separate_observables and bit_packed=False and transposed=True:
obs.dtype=bool_
obs.shape=(num_observables, shots)
Observable `m` in shot `s` is obs[m, s]
if separate_observables and bit_packed=True and transposed=False:
obs.dtype=uint8
obs.shape=(shots, math.ceil(num_observables / 8))
The bit for detection event `m` in shot `s` is at
(dets[s, m // 8] >> (m % 8)) & 1
The bit for observable `m` in shot `s` is at
(obs[s, m // 8] >> (m % 8)) & 1
obs.shape=(shots, (num_observables + 7) // 8)
Observable `m` in shot `s` is `(obs[s, m // 8] >> (m % 8)) & 1`
if separate_observables and bit_packed=True and transposed=True:
obs.dtype=uint8
obs.shape=(num_observables, (shots + 7) // 8)
Observable `m` in shot `s` is `(obs[m, s // 8] >> (s % 8)) & 1`

Examples:
>>> import stim
>>> c = stim.Circuit('''
... H 0
... CNOT 0 1
... R 0 1 2
... CX 0 1 2 1
... X_ERROR(1.0) 0
... M 0 1
... DETECTOR rec[-1] rec[-2]
... M 0 1 2
... DETECTOR rec[-2]
... DETECTOR rec[-1] rec[-2] rec[-3]
... OBSERVABLE_INCLUDE(0) rec[-1]
... ''')
>>> s = c.compile_detector_sampler()
>>> s.sample(shots=1)
array([[ True]])
>>> dets, obs = s.sample(shots=5, separate_observables=True)
>>> dets
array([[False, True],
[False, True],
[False, True],
[False, True],
[False, True]])
>>> obs
array([[False],
[False],
[False],
[False],
[False]])

>>> dets, obs = s.sample(
... shots=90,
... bit_packed=True,
... transposed=True,
... separate_observables=True,
... )
>>> dets
array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 3]],
dtype=uint8)
>>> obs
array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8)
"""
def sample_bit_packed(
self,
Expand Down
Loading
Loading