diff --git a/qiskit_machine_learning/connectors/torch_connector.py b/qiskit_machine_learning/connectors/torch_connector.py index 537d97d76..d975dc0bf 100644 --- a/qiskit_machine_learning/connectors/torch_connector.py +++ b/qiskit_machine_learning/connectors/torch_connector.py @@ -1,6 +1,6 @@ # This code is part of a Qiskit project. # -# (C) Copyright IBM 2021, 2023. +# (C) Copyright IBM 2021, 2024. # # This code is licensed under the Apache License, Version 2.0. You may # obtain a copy of this license in the LICENSE.txt file in the root directory @@ -13,7 +13,7 @@ """A connector to use Qiskit (Quantum) Neural Networks as PyTorch modules.""" from __future__ import annotations -from typing import Tuple, Any, cast +from typing import Tuple, Any, cast, Literal import numpy as np @@ -52,12 +52,57 @@ class Module: # type: ignore pass +def _get_einsum_signature(n_dimensions: int, return_type: Literal["input", "weight"]) -> str: + """ + Generate an Einstein summation signature for a given number of dimensions and return type. + + Args: + n_dimensions (int): The number of dimensions for the summation. + return_type (Literal["input", "weight"]): The type of the return signature. + - "input": Return signature includes all input indices except the last one. + - "weight": Return signature includes only the last index as the output. + + Returns: + str: The Einstein summation signature. + + Raises: + RuntimeError: If the number of dimensions exceeds the character limit. + ValueError: If an invalid return type is provided. + + Example: + Consider a scenario where n_dimensions is 3 and return_type is "input": + >>> _get_einsum_signature(3, "input") + 'ab,abc->ac' + This returns the Einstein summation signature 'ab,abc->ac' for input with three dimensions. + """ + trace = "" + char_limit = 26 + for i in range(n_dimensions): + trace += chr(97 + i) # chr(97) == 'a' + if i >= char_limit: + raise RuntimeError( + f"Cannot define an Einstein summation with more tha {char_limit:d} dimensions." + ) + + if return_type == "input": + signature = f"{trace[:-1]},{trace:s}->{trace[0] + trace[2:]}" + elif return_type == "weight": + signature = f"{trace[:-1]},{trace:s}->{trace[-1]}" + else: + raise ValueError( + f'The only allowed return types are ["input", "weight"], got {return_type:s} instead.' + ) + + return signature + + @_optionals.HAS_TORCH.require_in_instance class TorchConnector(Module): """Connects a Qiskit (Quantum) Neural Network to PyTorch.""" # pylint: disable=abstract-method class _TorchNNFunction(Function): + # pylint: disable=arguments-differ @staticmethod def forward( # type: ignore @@ -187,7 +232,9 @@ def backward(ctx: Any, grad_output: Tensor) -> Tuple: # type: ignore # able to do back-prop in a batched manner. # Pytorch does not support sparse einsum, so we rely on Sparse. # pylint: disable=no-member - input_grad = sparse.einsum("ij,ijk->ik", grad_coo, input_grad) + n_dimension = max(grad_coo.ndim, input_grad.ndim) + signature = _get_einsum_signature(n_dimension, return_type="input") + input_grad = sparse.einsum(signature, grad_coo, input_grad) # return sparse gradients input_grad = torch.sparse_coo_tensor(input_grad.coords, input_grad.data) @@ -205,7 +252,9 @@ def backward(ctx: Any, grad_output: Tensor) -> Tuple: # type: ignore input_grad = torch.as_tensor(input_grad, dtype=torch.float) # same as above - input_grad = torch.einsum("ij,ijk->ik", grad_output.detach().cpu(), input_grad) + n_dimension = max(grad_output.detach().cpu().ndim, input_grad.ndim) + signature = _get_einsum_signature(n_dimension, return_type="input") + input_grad = torch.einsum(signature, grad_output.detach().cpu(), input_grad) # place the resulting tensor to the device where they were stored input_grad = input_grad.to(input_data.device) @@ -226,7 +275,9 @@ def backward(ctx: Any, grad_output: Tensor) -> Tuple: # type: ignore # w.r.t. each parameter k. The weights' dimension is independent of the # batch size. # pylint: disable=no-member - weights_grad = sparse.einsum("ij,ijk->k", grad_coo, weights_grad) + n_dimension = max(grad_coo.ndim, weights_grad.ndim) + signature = _get_einsum_signature(n_dimension, return_type="weight") + weights_grad = sparse.einsum(signature, grad_coo, weights_grad) # return sparse gradients weights_grad = torch.sparse_coo_tensor( @@ -244,9 +295,9 @@ def backward(ctx: Any, grad_output: Tensor) -> Tuple: # type: ignore weights_grad = weights_grad.todense() weights_grad = torch.as_tensor(weights_grad, dtype=torch.float) # same as above - weights_grad = torch.einsum( - "ij,ijk->k", grad_output.detach().cpu(), weights_grad - ) + n_dimension = max(grad_output.detach().cpu().ndim, weights_grad.ndim) + signature = _get_einsum_signature(n_dimension, return_type="weight") + weights_grad = torch.einsum(signature, grad_output.detach().cpu(), weights_grad) # place the resulting tensor to the device where they were stored weights_grad = weights_grad.to(weights.device) diff --git a/releasenotes/notes/fix-716-mismatch-dimension-pytorch-ba01bea90eba1435.yaml b/releasenotes/notes/fix-716-mismatch-dimension-pytorch-ba01bea90eba1435.yaml new file mode 100644 index 000000000..c64e30f1d --- /dev/null +++ b/releasenotes/notes/fix-716-mismatch-dimension-pytorch-ba01bea90eba1435.yaml @@ -0,0 +1,6 @@ +--- +fixes: + - | + Fixes the dimension mismatch error in the `torch_connector` raised when using other-than 3D datasets. + The updated implementation defines the Einstein summation signature dynamically based on the number of + dimensions `ndim` of the input data (up to 26 dimensions). diff --git a/test/connectors/test_torch_connector.py b/test/connectors/test_torch_connector.py index eeb526dc5..57609e26b 100644 --- a/test/connectors/test_torch_connector.py +++ b/test/connectors/test_torch_connector.py @@ -1,6 +1,6 @@ # This code is part of a Qiskit project. # -# (C) Copyright IBM 2021, 2023. +# (C) Copyright IBM 2021, 2024. # # This code is licensed under the Apache License, Version 2.0. You may # obtain a copy of this license in the LICENSE.txt file in the root directory @@ -25,6 +25,7 @@ from qiskit_machine_learning import QiskitMachineLearningError from qiskit_machine_learning.connectors import TorchConnector from qiskit_machine_learning.neural_networks import SamplerQNN, EstimatorQNN +from qiskit_machine_learning.connectors.torch_connector import _get_einsum_signature @ddt @@ -44,6 +45,19 @@ def setup_test(self): torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]]), ] + def test_get_einsum_signature(self): + # Test valid inputs and outputs + self.assertEqual(_get_einsum_signature(3, "input"), "ab,abc->ac") + self.assertEqual(_get_einsum_signature(3, "weight"), "ab,abc->c") + + # Test raises for invalid return_type + with self.assertRaises(ValueError): + _get_einsum_signature(3, "invalid_type") + + # Test raises for exceeding character limit + with self.assertRaises(RuntimeError): + _get_einsum_signature(30, "input") + def _validate_backward_automatically(self, model: TorchConnector) -> None: """Uses PyTorch to validate the backward pass / autograd.