Improvement cuStabilizer DEM sampling backend and consolidate tests

[kvmto]

Summary

• Remove the torch fallback path from dem_sampling.py, making cuQuantum's cuStabilizer BitMatrixSampler the sole DEM sampling backend. This eliminates the USE_CUSTAB env-var toggle, the _use_custab() helper, and the custab_matrix_sampling() indirection — dem_sampling() now calls cuST directly.
• Simplify the sampler cache: replace the id(H) address-based cache key with a proper is-identity check (_cached_H is not H) and pre-cache H.T to avoid repeated transposes.
• Consolidate test files: merge test_dem_sampling_custab.py and test_dem_sampling_integration.py into test_dem_sampling.py, removing the now-unnecessary fallback-specific tests.
• Add cuquantum>=26.3.0 as an explicit dependency in requirements_public_train.txt.
• Fix CI to install from requirements_public_train.txt (which includes inference deps) instead of requirements_public_inference.txt alone.

Test plan

• CI GPU tests pass (test_dem_sampling.py covers CPU, GPU, statistical, integration, and cuST-vs-torch deterministic cross-checks)
• Verify cuquantum>=26.3.0 resolves correctly in the CI environment
• Confirm no regressions in training pipeline end-to-end

[ivanbasov]

CI failures root cause

Two related issues, both from the same design decision — making cuquantum.BitMatrixSampler the sole backend without adding skip guards in the tests.

---

Issue 1 — Unit-test CI (CPU runners): ModuleNotFoundError: No module named 'cuquantum'

dem_sampling() now does a hard import of cuquantum.stabilizer unconditionally at call time (line 81 of the new dem_sampling.py):

def dem_sampling(...):
    from cuquantum.stabilizer.dem_sampling import BitMatrixSampler  # <-- always runs
    from cuquantum.stabilizer.simulator import Options

The CPU CI runner doesn't have cuquantum installed, so every test that calls dem_sampling() crashes immediately. Affected test classes in test_dem_sampling.py:

• TestBitMatrixSamplerIsUsed — imports cuquantum directly, no skip guard
• TestDemSamplingCPU — calls dem_sampling(), no skip guard
• TestDEMSamplingIntegration — calls generate_batch() → dem_sampling(), no skip guard
• TestCustVsTorchDeterministic — calls dem_sampling(), no skip guard
• TestCustVsTorchStatistical.test_small_matrix_cpu — calls dem_sampling(), no skip guard

The old test_dem_sampling_custab.py handled this correctly with @unittest.skipUnless(_custab_available(), ...) at class level. That decorator was dropped when the files were consolidated.

---

Issue 2 — GPU CI runner: ModuleNotFoundError: No module named 'cuquantum.stabilizer'

The GPU runner has cuquantum installed, but an older version that predates the stabilizer subpackage (which requires >=26.3.0). Same test classes fail for the same reason — no skip guard catches the import error.

---

Suggested fix

The cleanest approach is to move the cuquantum import to module level with a try/except (restoring _CUSTAB_AVAILABLE), and add @unittest.skipUnless(_dem_mod._CUSTAB_AVAILABLE, "cuquantum>=26.3.0 not available") to all five affected test classes. The old code on main already does this correctly in _custab_available() / _CUSTAB_AVAILABLE — the pattern just needs to be carried forward.

# dem_sampling.py — module level
try:
    from cuquantum.stabilizer.dem_sampling import BitMatrixSampler
    from cuquantum.stabilizer.simulator import Options
    _CUSTAB_AVAILABLE = True
except ImportError:
    BitMatrixSampler = None
    Options = None
    _CUSTAB_AVAILABLE = False

# test_dem_sampling.py
@unittest.skipUnless(_dem_mod._CUSTAB_AVAILABLE, "cuquantum>=26.3.0 (stabilizer) not available")
class TestDemSamplingCPU(unittest.TestCase): ...

@unittest.skipUnless(_dem_mod._CUSTAB_AVAILABLE, "cuquantum>=26.3.0 (stabilizer) not available")
class TestBitMatrixSamplerIsUsed(unittest.TestCase): ...

# ... same for TestDEMSamplingIntegration, TestCustVsTorchDeterministic, TestCustVsTorchStatistical

This makes all five classes skip gracefully on CPU CI (no cuquantum) and on GPU CI until the runner is updated to cuquantum>=26.3.0, while still running correctly when the right version is present.

Otherwise, looks good!

[bmhowe23]

Note: CI passed with 3b8015d, so we could go with that. However, I prefer one commit later (ce055a2). It is currently failing the CI, but I think the cuquantum team is publishing a change that would fix that.

[bmhowe23]

Note: CI passed with 3b8015d, so we could go with that. However, I prefer one commit later (ce055a2). It is currently failing the CI, but I think the cuquantum team is publishing a change that would fix that.

The cuquantum issue won't have a published fix until April 10th, so I reverted ce055a2.

[bmhowe23]

Note: CI passed with 3b8015d, so we could go with that. However, I prefer one commit later (ce055a2). It is currently failing the CI, but I think the cuquantum team is publishing a change that would fix that.

The cuquantum issue won't have a published fix until April 10th, so I reverted ce055a2.

OK, they overachieved, and it is ready now: https://pypi.org/project/cuquantum-python/26.3.0.post0/

I would like to re-revert ce055a2 (i.e. revert 05e92f8). I will let the CI for Kevin's changes finish up before doing that though. This will be much cleaner to not have the cu12/cu13 split.

[bmhowe23]

Note: CI passed with 3b8015d, so we could go with that. However, I prefer one commit later (ce055a2). It is currently failing the CI, but I think the cuquantum team is publishing a change that would fix that.

The cuquantum issue won't have a published fix until April 10th, so I reverted ce055a2.

OK, they overachieved, and it is ready now: https://pypi.org/project/cuquantum-python/26.3.0.post0/

I would like to re-revert ce055a2 (i.e. revert 05e92f8). I will let the CI for Kevin's changes finish up before doing that though. This will be much cleaner to not have the cu12/cu13 split.

The CI is failing on cu12/cu13 auto-detection again, this time with a new error. It looks like this will have to wait for another day. I will revert the most recent commit (84c814b) and once it passes the CI, I propose we merge it.

[kvmto]

Note: CI passed with 3b8015d, so we could go with that. However, I prefer one commit later (ce055a2). It is currently failing the CI, but I think the cuquantum team is publishing a change that would fix that.

The cuquantum issue won't have a published fix until April 10th, so I reverted ce055a2.

OK, they overachieved, and it is ready now: https://pypi.org/project/cuquantum-python/26.3.0.post0/
I would like to re-revert ce055a2 (i.e. revert 05e92f8). I will let the CI for Kevin's changes finish up before doing that though. This will be much cleaner to not have the cu12/cu13 split.

The CI is failing on cu12/cu13 auto-detection again, this time with a new error. It looks like this will have to wait for another day. I will revert the most recent commit (84c814b) and once it passes the CI, I propose we merge it.

agree