From 93b24805f067a97d06d502cea13395985a8079b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?= Date: Mon, 16 Dec 2024 20:16:37 +0000 Subject: [PATCH] move the MPS shape chacker to Cpp --- .../lightning_tensor/tncuda/TNCuda.hpp | 25 ++++++ .../tncuda/bindings/LTensorTNCudaBindings.hpp | 13 +++ .../tncuda/tests/Tests_MPSTNCuda.cpp | 48 ++++++++++ pennylane_lightning/core/src/utils/Util.hpp | 30 +++++++ .../core/src/utils/tests/Test_Util.cpp | 12 +++ .../lightning_tensor/_tensornet.py | 71 +-------------- .../lightning_tensor/test_tensornet_class.py | 88 +------------------ 7 files changed, 131 insertions(+), 156 deletions(-) diff --git a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/TNCuda.hpp b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/TNCuda.hpp index f268af0474..e8a8ce5efa 100644 --- a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/TNCuda.hpp +++ b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/TNCuda.hpp @@ -252,6 +252,31 @@ class TNCuda : public TNCudaBase { host_data_size); } + /** + * @brief Check if the provided MPS has the correct dimension for C++ + * backend. + * + * @param MPS_shape_source Dimension list of incoming MPS. + */ + + void MPSShapeCheck( + const std::vector> &MPS_shape_source) { + bool sameShape = sitesExtents_ == MPS_shape_source; + if (!sameShape) { + auto MPS_shape_source_str = + Pennylane::Util::vector2DToString( + MPS_shape_source); + auto MPS_shape_dest_str = + Pennylane::Util::vector2DToString(sitesExtents_); + + PL_ABORT_IF_NOT( + sitesExtents_ == MPS_shape_source, + "The incoming MPS does not have the correct layout for " + "lightning.tensor.\n Incoming MPS: " + + MPS_shape_source_str + + "\n Destination MPS: " + MPS_shape_dest_str) + } + } /** * @brief Append multiple gates to the compute graph. * NOTE: This function does not update the quantum state but only appends diff --git a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/bindings/LTensorTNCudaBindings.hpp b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/bindings/LTensorTNCudaBindings.hpp index 1522305a82..ae4ae6fd66 100644 --- a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/bindings/LTensorTNCudaBindings.hpp +++ b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/bindings/LTensorTNCudaBindings.hpp @@ -137,6 +137,19 @@ void registerBackendClassSpecificBindingsMPS(PyClass &pyclass) { .def( "updateMPSSitesData", [](TensorNet &tensor_network, std::vector &tensors) { + // Extract the incoming MPS shape + std::vector> MPS_shape_source; + for (std::size_t idx = 0; idx < tensors.size(); idx++) { + py::buffer_info numpyArrayInfo = tensors[idx].request(); + auto MPS_site_source_shape = numpyArrayInfo.shape; + std::vector MPS_site_source( + MPS_site_source_shape.begin(), + MPS_site_source_shape.end()); + MPS_shape_source.push_back(std::move(MPS_site_source)); + } + + tensor_network.MPSShapeCheck(MPS_shape_source); + for (std::size_t idx = 0; idx < tensors.size(); idx++) { py::buffer_info numpyArrayInfo = tensors[idx].request(); auto *data_ptr = static_cast *>( diff --git a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/tests/Tests_MPSTNCuda.cpp b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/tests/Tests_MPSTNCuda.cpp index b75c272697..a359014bf4 100644 --- a/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/tests/Tests_MPSTNCuda.cpp +++ b/pennylane_lightning/core/src/simulators/lightning_tensor/tncuda/tests/Tests_MPSTNCuda.cpp @@ -120,6 +120,54 @@ TEMPLATE_TEST_CASE("MPSTNCuda::setIthMPSSite", "[MPSTNCuda]", float, double) { } } +TEMPLATE_TEST_CASE("MPSTNCuda::MPSShapeCheck()", "[MPSTNCuda]", float, double) { + const std::size_t num_qubits = 4; + const std::size_t maxBondDim = 8; + SECTION("Correct incoming MPS shape") { + + MPSTNCuda mps_state{num_qubits, maxBondDim}; + + std::vector> correct_shape{ + {2, 2}, {2, 2, 4}, {4, 2, 2}, {2, 2}}; + + REQUIRE_NOTHROW(mps_state.MPSShapeCheck(correct_shape)); + } + + SECTION("Incorrect incoming MPS shape, bond dimension") { + MPSTNCuda mps_state{num_qubits, maxBondDim}; + + std::vector> incorrect_shape{ + {2, 2}, {2, 2, 2}, {2, 2, 2}, {2, 2}}; + + REQUIRE_THROWS_WITH( + mps_state.MPSShapeCheck(incorrect_shape), + Catch::Matchers::Contains("The incoming MPS does not have the " + "correct layout for lightning.tensor")); + } + SECTION("Incorrect incoming MPS shape, physical dimension") { + MPSTNCuda mps_state{num_qubits, maxBondDim}; + + std::vector> incorrect_shape{ + {4, 2}, {2, 4, 4}, {4, 4, 2}, {2, 4}}; + + REQUIRE_THROWS_WITH( + mps_state.MPSShapeCheck(incorrect_shape), + Catch::Matchers::Contains("The incoming MPS does not have the " + "correct layout for lightning.tensor")); + } + SECTION("Incorrect incoming MPS shape, number sites") { + MPSTNCuda mps_state{num_qubits, maxBondDim}; + + std::vector> incorrect_shape{ + {2, 2}, {2, 2, 2}, {2, 2}}; + + REQUIRE_THROWS_WITH( + mps_state.MPSShapeCheck(incorrect_shape), + Catch::Matchers::Contains("The incoming MPS does not have the " + "correct layout for lightning.tensor")); + } +} + TEMPLATE_TEST_CASE("MPSTNCuda::SetBasisStates() & reset()", "[MPSTNCuda]", float, double) { std::vector> basisStates = { diff --git a/pennylane_lightning/core/src/utils/Util.hpp b/pennylane_lightning/core/src/utils/Util.hpp index 6f512ac40d..bf5d25651e 100644 --- a/pennylane_lightning/core/src/utils/Util.hpp +++ b/pennylane_lightning/core/src/utils/Util.hpp @@ -592,4 +592,34 @@ bool areVecsDisjoint(const std::vector &v1, const std::vector &v2) { } return true; } + +/** + * @brief Convert a 2D vector to string. + * @tparam T Data type. + * @param vec Vector to convert. + * + * @return std::string String with the vector values. + */ +template +std::string vector2DToString(const std::vector> &vec) { + std::ostringstream oss; + oss << "["; + + for (std::size_t i = 0; i < vec.size(); ++i) { + oss << "["; + for (std::size_t j = 0; j < vec[i].size(); ++j) { + oss << vec[i][j]; + if (j != vec[i].size() - 1) { + oss << ", "; // Add a comma between elements in the inner vector + } + } + oss << "]"; + if (i != vec.size() - 1) { + oss << ", "; // Add a comma between inner vectors + } + } + oss << "]"; + return oss.str(); // Return the resulting string +} + } // namespace Pennylane::Util diff --git a/pennylane_lightning/core/src/utils/tests/Test_Util.cpp b/pennylane_lightning/core/src/utils/tests/Test_Util.cpp index 2ba66c779f..709eb6dae8 100644 --- a/pennylane_lightning/core/src/utils/tests/Test_Util.cpp +++ b/pennylane_lightning/core/src/utils/tests/Test_Util.cpp @@ -233,3 +233,15 @@ TEST_CASE("Util::areVecsDisjoint", "[Util][LinearAlgebra]") { REQUIRE(areVecsDisjoint(vec0, vec1) == false); } } + +TEST_CASE("Utils::vector2DToString", "[Utils]") { + SECTION("Test for convert 2Dvector to string") { + std::vector> vec{ + {2, 2, 4}, {4, 2, 8}, {8, 2, 8}}; + std::string ref_str{"[[2, 2, 4], [4, 2, 8], [8, 2, 8]]"}; + + std::string vec2str = vector2DToString(vec); + + REQUIRE(ref_str == vec2str); + } +} \ No newline at end of file diff --git a/pennylane_lightning/lightning_tensor/_tensornet.py b/pennylane_lightning/lightning_tensor/_tensornet.py index 24f796f3e9..1684f40956 100644 --- a/pennylane_lightning/lightning_tensor/_tensornet.py +++ b/pennylane_lightning/lightning_tensor/_tensornet.py @@ -26,8 +26,6 @@ except ImportError: pass -from typing import List - import numpy as np import pennylane as qml from pennylane import BasisState, DeviceError, MPSPrep, StatePrep @@ -125,56 +123,6 @@ def gate_matrix_decompose(gate_ops_matrix, wires, max_mpo_bond_dim, c_dtype): return mpos, sorted_wires -def set_bond_dims(num_qubits: int, max_bond_dim: int) -> List: - """Compute the MPS bond dimensions on base to the number of wires.""" - - log_max_bond_dim = np.log2(max_bond_dim) - limit_dimension = 2 ** int(log_max_bond_dim) - localBondDims = [limit_dimension] * (num_qubits - 1) - - for i in range(len(localBondDims)): - bondDim = min(i + 1, num_qubits - i - 1) - if bondDim <= log_max_bond_dim: - localBondDims[i] = 2**bondDim - - return localBondDims - - -def set_sites_extents(num_qubits: int, max_bond_dim: int) -> List: - """Compute the MPS sites dimensions on base to the number of wires.""" - - bondDims = set_bond_dims(num_qubits, max_bond_dim) - qubitDims = [2 for _ in range(num_qubits)] - - localSiteExtents = [] - for i in range(num_qubits): - if i == 0: - localSite = [qubitDims[i], bondDims[i]] - elif i == num_qubits - 1: - localSite = [bondDims[i - 1], qubitDims[i]] - else: - localSite = [bondDims[i - 1], qubitDims[i], bondDims[i]] - - localSiteExtents.append(localSite) - - return localSiteExtents - - -def MPSPrep_check(MPS: List, num_wires: int, max_bond_dim: int) -> None: - """Check if the provided MPS has the correct dimension for C++ backend.""" - - MPS_shape_dest = set_sites_extents(num_wires, max_bond_dim) - - MPS_shape_source = [list(site.shape) for site in MPS] - - same_shape = [s == d for s, d in zip(MPS_shape_source, MPS_shape_dest)] - - if not all(same_shape): - raise ValueError( - f"The custom MPS does not have the correct layout for lightning.tensor.\n MPS source shape {MPS_shape_source}\n MPS destination shape {MPS_shape_dest}" - ) - - # pylint: disable=too-many-instance-attributes class LightningTensorNet: """Lightning tensornet class. @@ -376,22 +324,6 @@ def _apply_basis_state(self, state, wires): self._tensornet.setBasisState(state) - def _load_mps_state(self, state: List): - """Prepares an initial state using MPS. - - Args: - state (List): A list of different numpy array with the MPS sites values. The structure should be as follows: - [ (2, 2), (2, 2, 4), (4, 2, 8), ..., - (8, 2, 4), (4, 2, 2), (2, 2) ] - wires (List): wires that the provided computational state should be - initialized on. - - Note: The correct MPS sites format and layout are user responsible. - """ - mps = state.mps - MPSPrep_check(mps, self._num_wires, self._max_bond_dim) - self._tensornet.updateMPSSitesData(mps) - def _apply_MPO(self, gate_matrix, wires): """Apply a matrix product operator to the quantum state (MPS method only). @@ -513,7 +445,8 @@ def apply_operations(self, operations): self._apply_basis_state(operations[0].parameters[0], operations[0].wires) operations = operations[1:] elif isinstance(operations[0], MPSPrep): - self._load_mps_state(operations[0]) + mps = operations[0].mps + self._tensornet.updateMPSSitesData(mps) operations = operations[1:] self._apply_lightning(operations) diff --git a/tests/lightning_tensor/test_tensornet_class.py b/tests/lightning_tensor/test_tensornet_class.py index 74530c47d7..64f367ed7a 100644 --- a/tests/lightning_tensor/test_tensornet_class.py +++ b/tests/lightning_tensor/test_tensornet_class.py @@ -25,11 +25,8 @@ else: from pennylane_lightning.lightning_tensor._tensornet import ( LightningTensorNet, - MPSPrep_check, decompose_dense, gate_matrix_decompose, - set_bond_dims, - set_sites_extents, ) if not LightningDevice._CPP_BINARY_AVAILABLE: # pylint: disable=protected-access @@ -74,7 +71,7 @@ def test_wrong_device_name(): def test_wrong_method_name(): - """Test an invalid device name""" + """Test an invalid method name""" with pytest.raises(qml.DeviceError, match="The method "): LightningTensorNet(3, max_bond_dim=5, device_name="lightning.tensor", method="spider_web") @@ -138,86 +135,3 @@ def test_gate_matrix_decompose(): assert np.allclose(sorted_wired, sorted(wires), atol=1e-6) assert np.allclose(unitary_f, original_gate, atol=1e-6) - - -@pytest.mark.parametrize( - "n_qubits,max_bond,expected", - [ - (2, 128, [2]), - (8, 128, [2, 4, 8, 16, 8, 4, 2]), - (8, 8, [2, 4, 8, 8, 8, 4, 2]), - (15, 2, [2 for _ in range(14)]), - ], -) -def test_set_bond_dims(n_qubits, max_bond, expected): - - result = set_bond_dims(n_qubits, max_bond) - - assert len(result) == len(expected) - assert all([a == b for a, b in zip(result, expected)]) - - -@pytest.mark.parametrize( - "n_qubits,max_bond,expected", - [ - (2, 128, [[2, 2], [2, 2]]), - ( - 8, - 128, - [[2, 2], [2, 2, 4], [4, 2, 8], [8, 2, 16], [16, 2, 8], [8, 2, 4], [4, 2, 2], [2, 2]], - ), - (8, 8, [[2, 2], [2, 2, 4], [4, 2, 8], [8, 2, 8], [8, 2, 8], [8, 2, 4], [4, 2, 2], [2, 2]]), - (15, 2, [[2, 2]] + [[2, 2, 2] for _ in range(13)] + [[2, 2]]), - ], -) -def test_set_sites_extents(n_qubits, max_bond, expected): - - result = set_sites_extents(n_qubits, max_bond) - - assert len(result) == len(expected) - assert all([a == b for a, b in zip(result, expected)]) - - -@pytest.mark.parametrize( - "wires,max_bond,MPS_shape", - [ - (2, 128, [[2, 2], [2, 2]]), - ( - 8, - 128, - [[2, 2], [2, 2, 4], [4, 2, 8], [8, 2, 16], [16, 2, 8], [8, 2, 4], [4, 2, 2], [2, 2]], - ), - (8, 8, [[2, 2], [2, 2, 4], [4, 2, 8], [8, 2, 8], [8, 2, 8], [8, 2, 4], [4, 2, 2], [2, 2]]), - (15, 2, [[2, 2]] + [[2, 2, 2] for _ in range(13)] + [[2, 2]]), - ], -) -def test_MPSPrep_check_pass(wires, max_bond, MPS_shape): - MPS = [np.zeros(i) for i in MPS_shape] - - MPSPrep_check(MPS, wires, max_bond) - - -@pytest.mark.parametrize( - "wires,max_bond,MPS_shape", - [ - (2, 128, [[2, 3], [3, 2]]), # Incorrect bond dim. - ( - 8, - 128, - [[2, 2], [2, 4, 4], [4, 4, 8], [8, 4, 16], [16, 4, 8], [8, 4, 4], [4, 4, 2], [2, 2]], - ), # Incorrect physical dim. - ( - 8, - 8, - [[2, 2], [3, 2, 4], [4, 2, 8], [8, 2, 8], [8, 2, 8], [8, 2, 4], [4, 2, 2], [2, 2]], - ), # Incorrect only one bond dim. - (15, 2, [[2, 2]] + [[2, 2, 2] for _ in range(14)] + [[2, 2]]), # Incorrect amount of sites - ], -) -def test_MPSPrep_check_fail(wires, max_bond, MPS_shape): - MPS = [np.zeros(i) for i in MPS_shape] - - with pytest.raises( - ValueError, match="The custom MPS does not have the correct layout for lightning.tensor" - ): - MPSPrep_check(MPS, wires, max_bond)