[WIP] SamplerV2

qiskit_ibm_runtime/base_primitive.py

@@ -127,7 +127,7 @@ def _run(self, pubs: Union[list[EstimatorPub], list[SamplerPub]]) -> RuntimeJob:
         Returns:
             Submitted job.
         """
-        primitive_inputs = {"tasks": pubs}
+        primitive_inputs = {"pubs": pubs}
         options_dict = asdict(self.options)
         self._validate_options(options_dict)
         primitive_inputs.update(self._options_class._get_program_inputs(options_dict))

qiskit_ibm_runtime/qiskit/primitives/sampler_pub.py

@@ -1,6 +1,6 @@
 # This code is part of Qiskit.
 #
-# (C) Copyright IBM 2023.
+# (C) Copyright IBM 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
@@ -9,68 +9,130 @@
 # Any modifications or derivative works of this code must retain this
 # copyright notice, and modified files need to carry a notice indicating
 # that they have been altered from the originals.
-# type: ignore
+
 
 """
-Sampler PUB class
+Sampler Pub class
 """
 
 from __future__ import annotations
 
-from dataclasses import dataclass, field
-from typing import Union, Optional, Tuple
+from typing import Tuple, Union
+from numbers import Integral
 
 from qiskit import QuantumCircuit
 
-from .base_pub import BasePub
 from .bindings_array import BindingsArray, BindingsArrayLike
 from .shape import ShapedMixin
 
 
-@dataclass(frozen=True)
-class SamplerPub(BasePub, ShapedMixin):
-    """PUB for Sampler.
-    PUB is composed of double (circuit, parameter_values).
+class SamplerPub(ShapedMixin):
+    """Pub (Primitive Unified Bloc) for Sampler.
+
+    Pub is composed of tuple (circuit, parameter_values, shots).
+
+    If shots are provided this number of shots will be run with the sampler,
+    if ``shots=None`` the number of run shots is determined by the sampler.
     """
 
-    parameter_values: Optional[BindingsArray] = BindingsArray([], shape=())
-    _shape: tuple[int, ...] = field(init=False)
+    def __init__(
+        self,
+        circuit: QuantumCircuit,
+        parameter_values: BindingsArray | None = None,
+        shots: int | None = None,
+        validate: bool = True,
+    ):
+        """Initialize a sampler pub.
 
-    def __post_init__(self):
-        super().__setattr__("_shape", self.parameter_values.shape)
+        Args:
+            circuit: A quantum circuit.
+            parameter_values: A bindings array.
+            shots: A specific number of shots to run with. This value takes
+                precedence over any value owed by or supplied to a sampler.
+            validate: If ``True``, the input data is validated during initialization.
+        """
+        super().__init__()
+        self._circuit = circuit
+        self._parameter_values = parameter_values or BindingsArray()
+        self._shots = shots
+        self._shape = self._parameter_values.shape
+        if validate:
+            self.validate()
+
+    @property
+    def circuit(self) -> QuantumCircuit:
+        """A quantum circuit."""
+        return self._circuit
+
+    @property
+    def parameter_values(self) -> BindingsArray:
+        """A bindings array."""
+        return self._parameter_values
+
+    @property
+    def shots(self) -> int | None:
+        """An specific number of shots to run with (optional).
+
+        This value takes precedence over any value owed by or supplied to a sampler.
+        """
+        return self._shots
 
     @classmethod
-    def coerce(cls, pub: SamplerPubLike) -> SamplerPub:
-        """Coerce SamplerPubLike into SamplerPub.
+    def coerce(cls, pub: SamplerPubLike, shots: int | None = None) -> SamplerPub:
+        """Coerce a :class:`~.SamplerPubLike` object into a :class:`~.SamplerPub` instance.
 
         Args:
-            pub: an object to be sampler pub.
+            pub: An object to coerce.
+            shots: An optional default number of shots to use if not
+                   already specified by the pub-like object.
 
         Returns:
-            A coerced estiamtor pub.
-
-        Raises:
-            ValueError: If input values are invalid.
+            A coerced sampler pub.
         """
+        # Validate shots kwarg if provided
+        if shots is not None:
+            if not isinstance(shots, Integral) or isinstance(shots, bool):
+                raise TypeError("shots must be an integer")
+            if shots < 0:
+                raise ValueError("shots must be non-negative")
+
         if isinstance(pub, SamplerPub):
+            if pub.shots is None and shots is not None:
+                return cls(
+                    circuit=pub.circuit,
+                    parameter_values=pub.parameter_values,
+                    shots=shots,
+                    validate=False,  # Assume Pub is already validated
+                )
             return pub
-        if len(pub) != 1 and len(pub) != 2:
-            raise ValueError(f"The length of pub must be 1 or 2, but length {len(pub)} is given.")
+
+        if isinstance(pub, QuantumCircuit):
+            return cls(circuit=pub, shots=shots, validate=True)
+
+        if len(pub) not in [1, 2, 3]:
+            raise ValueError(
+                f"The length of pub must be 1, 2 or 3, but length {len(pub)} is given."
+            )
         circuit = pub[0]
-        parameter_values = (
-            BindingsArray.coerce(pub[1]) if len(pub) == 2 else BindingsArray([], shape=())
-        )
-        return cls(circuit=circuit, parameter_values=parameter_values)
+        parameter_values = BindingsArray.coerce(pub[1]) if len(pub) > 1 else None
+        if len(pub) > 2 and pub[2] is not None:
+            shots = pub[2]
+        return cls(circuit=circuit, parameter_values=parameter_values, shots=shots, validate=True)
 
-    def validate(self) -> None:
-        """Validate the pub.
+    def validate(self):
+        """Validate the pub."""
+        if not isinstance(self.circuit, QuantumCircuit):
+            raise TypeError("circuit must be QuantumCircuit.")
 
-        Raises:
-            ValueError: If input values are invalid.
-        """
-        super().validate()
         self.parameter_values.validate()
-        # Cross validate circuits and paramter_values
+
+        if self.shots is not None:
+            if not isinstance(self.shots, Integral) or isinstance(self.shots, bool):
+                raise TypeError("shots must be an integer")
+            if self.shots < 0:
+                raise ValueError("shots must be non-negative")
+
+        # Cross validate circuits and parameter values
         num_parameters = self.parameter_values.num_parameters
         if num_parameters != self.circuit.num_parameters:
             raise ValueError(
@@ -79,4 +141,10 @@ def validate(self) -> None:
             )
 
 
-SamplerPubLike = Union[SamplerPub, Tuple[QuantumCircuit, BindingsArrayLike]]
+SamplerPubLike = Union[
+    SamplerPub,
+    QuantumCircuit,
+    Tuple[QuantumCircuit],
+    Tuple[QuantumCircuit, BindingsArrayLike],
+    Tuple[QuantumCircuit, BindingsArrayLike, Union[Integral, None]],
+]

qiskit_ibm_runtime/sampler.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 import os
-from typing import Dict, Optional, Sequence, Any, Union
+from typing import Dict, Optional, Sequence, Any, Union, Iterable
 import logging
 
 from qiskit.circuit import QuantumCircuit
@@ -32,6 +32,7 @@
 
 # TODO: remove when we have real v2 base estimator
 from .qiskit.primitives import BaseSamplerV2
+from .qiskit.primitives.sampler_pub import SamplerPubLike, SamplerPub
 
 logger = logging.getLogger(__name__)
 
@@ -90,11 +91,36 @@ def __init__(
         Sampler.__init__(self)
         BasePrimitiveV2.__init__(self, backend=backend, session=session, options=options)
 
-        raise NotImplementedError("SamplerV2 is not currently supported.")
-
         # if self._service._channel_strategy == "q-ctrl":
         #     raise NotImplementedError("SamplerV2 is not supported with q-ctrl channel strategy.")
 
+    def run(
+            self, pubs: SamplerPubLike | Iterable[SamplerPubLike]
+    ) -> RuntimeJob:
+        """Submit a request to the sampler primitive.
+
+        Args:
+            pubs: A Primitive Unified Bloc or a list of such blocs
+
+        Returns:
+            Submitted job.
+            The result of the job is an instance of :class:`qiskit.primitives.containers.PrimitiveResult`.
+
+        Raises:
+            ValueError: Invalid arguments are given.
+        """
+        if isinstance(pubs, SamplerPub):
+            pubs = [pubs]
+        elif isinstance(pubs, tuple) and isinstance(pubs[0], QuantumCircuit):
+            pubs = [SamplerPub.coerce(pubs)]
+        elif pubs is not SamplerPub:
+            pubs = [SamplerPub.coerce(pub) for pub in pubs]
+
+        for pub in pubs:
+            pub.validate()
+
+        return self._run(pubs)
+
     def _validate_options(self, options: dict) -> None:
         """Validate that program inputs (options) are valid
 

qiskit_ibm_runtime/utils/json.py

@@ -70,6 +70,7 @@
 # TODO: Remove when they are in terra
 from ..qiskit.primitives import ObservablesArray, BindingsArray
 from ..qiskit.primitives.base_pub import BasePub
+from ..qiskit.primitives.sampler_pub import SamplerPub
 
 _TERRA_VERSION = tuple(
     int(x) for x in re.match(r"\d+\.\d+\.\d", _terra_version_string).group(0).split(".")[:3]
@@ -256,6 +257,8 @@ def default(self, obj: Any) -> Any:  # pylint: disable=arguments-differ
             return {"__type__": "Instruction", "__value__": value}
         if isinstance(obj, BasePub):
             return asdict(obj)
+        if isinstance(obj, SamplerPub):
+            return {'circuit': obj.circuit, 'parameter_values': obj.parameter_values, 'shots': obj.shots, 'shape': obj.shape}
         if isinstance(obj, ObservablesArray):
             return obj.tolist()
         if isinstance(obj, BindingsArray):