Add shot measurement support to lightning.tensor

.github/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* Add shot measurement support to `lightning.tensor`.
+  [(#852)](https://github.com/PennyLaneAI/pennylane-lightning/pull/852)
+
 * Build and upload Lightning-Tensor wheels (x86_64, AARCH64) to PyPI.
   [(#862)](https://github.com/PennyLaneAI/pennylane-lightning/pull/862)
 

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.39.0-dev17"
+__version__ = "0.39.0-dev18"

pennylane_lightning/core/src/bindings/Bindings.hpp

@@ -746,7 +746,26 @@ void registerLightningTensorBackendAgnosticMeasurements(PyClass &pyclass) {
             [](MeasurementsT &M, const std::shared_ptr<ObservableT> &ob) {
                 return M.var(*ob);
             },
-            "Variance of an observable object.");
+            "Variance of an observable object.")
+        .def("generate_samples", [](MeasurementsT &M,
+                                    const std::vector<std::size_t> &wires,
+                                    const std::size_t num_shots) {
+            const std::size_t num_wires = wires.size();
+            auto &&result = M.generate_samples(wires, num_shots);
+            const std::size_t ndim = 2;
+            const std::vector<std::size_t> shape{num_shots, num_wires};
+            constexpr auto sz = sizeof(std::size_t);
+            const std::vector<std::size_t> strides{sz * num_wires, sz};
+            // return 2-D NumPy array
+            return py::array(py::buffer_info(
+                result.data(), /* data as contiguous array  */
+                sz,            /* size of one scalar        */
+                py::format_descriptor<std::size_t>::format(), /* data type */
+                ndim,   /* number of dimensions      */
+                shape,  /* shape of the matrix       */
+                strides /* strides for each axis     */
+                ));
+        });
 }
 
 /**

pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/measurements/MeasurementsTNCuda.hpp

@@ -20,6 +20,7 @@
 
 #pragma once
 
+#include <algorithm>
 #include <complex>
 #include <cuComplex.h>
 #include <cutensornet.h>
@@ -164,6 +165,97 @@ template <class TensorNetT> class MeasurementsTNCuda {
         return h_res;
     }
 
+    /**
+     * @brief Utility method for samples.
+     *
+     * @param wires Wires can be a subset or the full system.
+     * @param num_samples Number of samples
+     * @param numHyperSamples Number of hyper samples to use in the calculation
+     * and is default as 1.
+     *
+     * @return std::vector<std::size_t> A 1-d array storing the samples.
+     * Each sample has a length equal to the number of wires. Each sample can
+     * be accessed using the stride `sample_id * num_wires`, where `sample_id`
+     * is a number between `0` and `num_samples - 1`.
+     */
+    auto generate_samples(const std::vector<std::size_t> &wires,
+                          const std::size_t num_samples,
+                          const int32_t numHyperSamples = 1)
+        -> std::vector<std::size_t> {
+        std::vector<int64_t> samples(num_samples * wires.size());
+
+        const std::vector<int32_t> modesToSample =
+            cuUtil::NormalizeCastIndices<std::size_t, int32_t>(
+                wires, tensor_network_.getNumQubits());
+
+        cutensornetStateSampler_t sampler;
+
+        PL_CUTENSORNET_IS_SUCCESS(cutensornetCreateSampler(
+            /* const cutensornetHandle_t */ tensor_network_.getTNCudaHandle(),
+            /* cutensornetState_t */ tensor_network_.getQuantumState(),
+            /* int32_t numModesToSample */ modesToSample.size(),
+            /* const int32_t *modesToSample */ modesToSample.data(),
+            /* cutensornetStateSampler_t * */ &sampler));
+
+        // Configure the quantum circuit sampler
+        const cutensornetSamplerAttributes_t samplerAttributes =
+            CUTENSORNET_SAMPLER_CONFIG_NUM_HYPER_SAMPLES;
+
+        PL_CUTENSORNET_IS_SUCCESS(cutensornetSamplerConfigure(
+            /* const cutensornetHandle_t */ tensor_network_.getTNCudaHandle(),
+            /* cutensornetStateSampler_t */ sampler,
+            /* cutensornetSamplerAttributes_t */ samplerAttributes,
+            /* const void *attributeValue */ &numHyperSamples,
+            /* size_t attributeSize */ sizeof(numHyperSamples)));
+
+        cutensornetWorkspaceDescriptor_t workDesc;
+        PL_CUTENSORNET_IS_SUCCESS(cutensornetCreateWorkspaceDescriptor(
+            /* const cutensornetHandle_t */ tensor_network_.getTNCudaHandle(),
+            /* cutensornetWorkspaceDescriptor_t * */ &workDesc));
+
+        const std::size_t scratchSize = cuUtil::getFreeMemorySize() / 2;
+
+        // Prepare the quantum circuit sampler for sampling
+        PL_CUTENSORNET_IS_SUCCESS(cutensornetSamplerPrepare(
+            /* const cutensornetHandle_t */ tensor_network_.getTNCudaHandle(),
+            /* cutensornetStateSampler_t */ sampler,
+            /* size_t maxWorkspaceSizeDevice */ scratchSize,
+            /* cutensornetWorkspaceDescriptor_t */ workDesc,
+            /* cudaStream_t unused as of v24.08 */ 0x0));
+
+        std::size_t worksize =
+            getWorkSpaceMemorySize(tensor_network_.getTNCudaHandle(), workDesc);
+
+        PL_ABORT_IF(worksize > scratchSize,
+                    "Insufficient workspace size on Device.\n");
+
+        const std::size_t d_scratch_length = worksize / sizeof(size_t) + 1;
+        DataBuffer<std::size_t> d_scratch(d_scratch_length,
+                                          tensor_network_.getDevTag(), true);
+
+        setWorkSpaceMemory(tensor_network_.getTNCudaHandle(), workDesc,
+                           reinterpret_cast<void *>(d_scratch.getData()),
+                           worksize);
+
+        PL_CUTENSORNET_IS_SUCCESS(cutensornetSamplerSample(
+            /* const cutensornetHandle_t */ tensor_network_.getTNCudaHandle(),
+            /* cutensornetStateSampler_t */ sampler,
+            /* int64_t numShots */ num_samples,
+            /* cutensornetWorkspaceDescriptor_t */ workDesc,
+            /* int64_t * */ samples.data(),
+            /* cudaStream_t unused as of v24.08  */ 0x0));
+
+        PL_CUTENSORNET_IS_SUCCESS(
+            cutensornetDestroyWorkspaceDescriptor(workDesc));
+        PL_CUTENSORNET_IS_SUCCESS(cutensornetDestroySampler(sampler));
+
+        std::vector<std::size_t> samples_size_t(samples.size());
+
+        std::transform(samples.begin(), samples.end(), samples_size_t.begin(),
+                       [](int64_t x) { return static_cast<std::size_t>(x); });
+        return samples_size_t;
+    }
+
     /**
      * @brief Calculate var value for a general ObservableTNCuda Observable.
      *

pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/measurements/tests/Test_MPSTNCuda_Measure.cpp

@@ -26,13 +26,15 @@
 #include "MPSTNCuda.hpp"
 #include "MeasurementsTNCuda.hpp"
 #include "TNCudaGateCache.hpp"
+#include "TestHelpers.hpp"
 #include "cuda_helpers.hpp"
 
 /// @cond DEV
 namespace {
 using namespace Pennylane::LightningTensor::TNCuda::Measures;
 using namespace Pennylane::LightningTensor::TNCuda::Observables;
 using namespace Pennylane::LightningTensor::TNCuda;
+using namespace Pennylane::Util;
 } // namespace
 /// @endcond
 
@@ -92,3 +94,47 @@ TEMPLATE_TEST_CASE("Probabilities", "[Measures]", float, double) {
         REQUIRE_THROWS_AS(measure.probs({2, 1}), LightningException);
     }
 }
+
+TEMPLATE_TEST_CASE("Samples", "[Measures]", float, double) {
+    using TensorNetT = MPSTNCuda<TestType>;
+
+    SECTION("Looping over different wire configurations:") {
+        // Probabilities calculated with Pennylane default.qubit:
+        std::vector<TestType> expected_probabilities = {
+            0.67078706, 0.03062806, 0.0870997,  0.00397696,
+            0.17564072, 0.00801973, 0.02280642, 0.00104134};
+
+        // Defining the State Vector that will be measured.
+        std::size_t bondDim = GENERATE(4, 5);
+        std::size_t num_qubits = 3;
+        std::size_t maxBondDim = bondDim;
+
+        TensorNetT mps_state{num_qubits, maxBondDim};
+
+        mps_state.applyOperations(
+            {{"RX"}, {"RX"}, {"RY"}, {"RY"}, {"RX"}, {"RY"}},
+            {{0}, {0}, {1}, {1}, {2}, {2}},
+            {{false}, {false}, {false}, {false}, {false}, {false}},
+            {{0.5}, {0.5}, {0.2}, {0.2}, {0.5}, {0.5}});
+        mps_state.append_mps_final_state();
+
+        auto measure = MeasurementsTNCuda<TensorNetT>(mps_state);
+
+        std::size_t num_samples = 100000;
+        const std::vector<std::size_t> wires = {0, 1, 2};
+        auto samples = measure.generate_samples(wires, num_samples);
+        auto counts = samples_to_decimal(samples, num_qubits, num_samples);
+
+        // compute estimated probabilities from histogram
+        std::vector<TestType> probabilities(counts.size());
+        for (std::size_t i = 0; i < counts.size(); i++) {
+            probabilities[i] = counts[i] / static_cast<TestType>(num_samples);
+        }
+
+        // compare estimated probabilities to real probabilities
+        SECTION("No wires provided:") {
+            REQUIRE_THAT(probabilities,
+                         Catch::Approx(expected_probabilities).margin(.1));
+        }
+    }
+}

pennylane_lightning/core/src/utils/TestHelpers.hpp

@@ -598,6 +598,35 @@ auto randomUnitary(RandomEngine &re, std::size_t num_qubits)
     return res;
 }
 
+inline auto samples_to_decimal(const std::vector<std::size_t> &samples,
+                               const std::size_t num_qubits,
+                               const std::size_t num_samples)
+    -> std::vector<std::size_t> {
+    constexpr uint32_t twos[] = {
+        1U << 0U,  1U << 1U,  1U << 2U,  1U << 3U,  1U << 4U,  1U << 5U,
+        1U << 6U,  1U << 7U,  1U << 8U,  1U << 9U,  1U << 10U, 1U << 11U,
+        1U << 12U, 1U << 13U, 1U << 14U, 1U << 15U, 1U << 16U, 1U << 17U,
+        1U << 18U, 1U << 19U, 1U << 20U, 1U << 21U, 1U << 22U, 1U << 23U,
+        1U << 24U, 1U << 25U, 1U << 26U, 1U << 27U, 1U << 28U, 1U << 29U,
+        1U << 30U, 1U << 31U};
+
+    std::size_t N = std::pow(2, num_qubits);
+    std::vector<std::size_t> counts(N, 0);
+    std::vector<std::size_t> samples_decimal(num_samples, 0);
+
+    // convert samples to decimal and then bin them in counts
+    for (std::size_t i = 0; i < num_samples; i++) {
+        for (std::size_t j = 0; j < num_qubits; j++) {
+            if (samples[i * num_qubits + j] != 0) {
+                samples_decimal[i] += twos[(num_qubits - 1 - j)];
+            }
+        }
+        counts[samples_decimal[i]] += 1;
+    }
+
+    return counts;
+}
+
 #define PL_REQUIRE_THROWS_MATCHES(expr, type, message_match)                   \
     REQUIRE_THROWS_AS(expr, type);                                             \
     REQUIRE_THROWS_WITH(expr, Catch::Matchers::Contains(message_match));

pennylane_lightning/core/src/utils/cuda_utils/tests/CMakeLists.txt

@@ -1,6 +1,6 @@
 cmake_minimum_required(VERSION 3.20)
 
-project(lightning_gpu_utils_tests)
+project(cuda_utils_tests)
 
 # Default build type for test code is Debug
 if(NOT CMAKE_BUILD_TYPE)