Return DShapeArray objects in measure primitives when possible

doc/releases/changelog-dev.md

@@ -293,6 +293,9 @@
   this nested module into other IRs/binary formats and lowering `call_function_in_module`
   to something that can dispatch calls to another runtime / VM.
 
+* Measurement primitives now support dynamic shape at the frontend, although at the PennyLane
+  side, the corresponding operations still lack such support.
+  [(#1170)](https://github.com/PennyLaneAI/catalyst/pull/1170)
 
 <h3>Contributors</h3>
 

frontend/catalyst/api_extensions/differentiation.py

@@ -31,7 +31,7 @@
 from pennylane import QNode
 
 import catalyst
-from catalyst.jax_extras import Jaxpr
+from catalyst.jax_extras import Jaxpr, make_jaxpr2
 from catalyst.jax_primitives import (
     GradParams,
     expval_p,
@@ -800,8 +800,7 @@ def _make_jaxpr_check_differentiable(
     return the output tree."""
     method = grad_params.method
     with mark_gradient_tracing(method):
-        jaxpr, shape = jax.make_jaxpr(f, return_shape=True)(*args, **kwargs)
-    _, out_tree = tree_flatten(shape)
+        jaxpr, _, out_tree = make_jaxpr2(f)(*args, **kwargs)
 
     for pos, arg in enumerate(jaxpr.in_avals):
         if arg.dtype.kind != "f" and pos in grad_params.expanded_argnums:

frontend/catalyst/jax_primitives.py

@@ -1633,10 +1633,8 @@ def _hamiltonian_lowering(jax_ctx: mlir.LoweringRuleContext, coeffs: ir.Value, *
 def _sample_abstract_eval(obs, shots, shape):
     assert isinstance(obs, AbstractObs)
 
-    if obs.primitive is compbasis_p:
-        assert shape == (shots, obs.num_qubits)
-    else:
-        assert shape == (shots,)
+    if Signature.is_dynamic_shape(shape):
+        return core.DShapedArray(shape, np.dtype("float64"))
 
     return core.ShapedArray(shape, jax.numpy.float64)
 
@@ -1670,10 +1668,10 @@ def _counts_def_impl(ctx, obs, shots, shape):  # pragma: no cover
 def _counts_abstract_eval(obs, shots, shape):
     assert isinstance(obs, AbstractObs)
 
-    if obs.primitive is compbasis_p:
-        assert shape == (2**obs.num_qubits,)
-    else:
-        assert shape == (2,)
+    if Signature.is_dynamic_shape(shape):
+        return core.DShapedArray(shape, np.dtype("float64")), core.DShapedArray(
+            shape, np.dtype("int64")
+        )
 
     return core.ShapedArray(shape, jax.numpy.float64), core.ShapedArray(shape, jax.numpy.int64)
 
@@ -1762,10 +1760,10 @@ def _var_lowering(jax_ctx: mlir.LoweringRuleContext, obs: ir.Value, shots: int,
 def _probs_abstract_eval(obs, shape, shots=None):
     assert isinstance(obs, AbstractObs)
 
-    if obs.primitive is compbasis_p:
-        assert shape == (2**obs.num_qubits,)
-    else:
-        raise TypeError("probs only supports computational basis")
+    assert obs.primitive is compbasis_p, "probs only supports computational basis"
+
+    if Signature.is_dynamic_shape(shape):
+        return core.DShapedArray(shape, np.dtype("float64"))
 
     return core.ShapedArray(shape, jax.numpy.float64)
 
@@ -1791,10 +1789,10 @@ def _probs_lowering(jax_ctx: mlir.LoweringRuleContext, obs: ir.Value, shape: tup
 def _state_abstract_eval(obs, shape, shots=None):
     assert isinstance(obs, AbstractObs)
 
-    if obs.primitive is compbasis_p:
-        assert shape == (2**obs.num_qubits,)
-    else:
-        raise TypeError("state only supports computational basis")
+    assert obs.primitive is compbasis_p, "state only supports computational basis"
+
+    if Signature.is_dynamic_shape(shape):
+        return core.DShapedArray(shape, np.dtype("complex128"))
 
     return core.ShapedArray(shape, jax.numpy.complex128)
 

frontend/catalyst/utils/calculate_grad_shape.py

@@ -17,7 +17,7 @@
 """
 
 
-from jax.core import ShapedArray
+from jax.core import DShapedArray, ShapedArray, Tracer
 
 
 class Signature:
@@ -76,15 +76,33 @@ def get_results(self):
 
     @staticmethod
     def is_tensor(x):
-        """Determine whether a type ``x`` is a ``jax.core.ShapedArray``.
+        """Determine whether a type ``x`` is a ``jax.core.DShapedArray``
+        or ``jax.core.ShapedArray``.
 
         Args:
             x: The type to be tested.
 
         Returns:
-            bool: Whether the type ``x`` is a ``jax.core.ShapedArray``
+            bool: Whether the type ``x`` is a ``jax.core.DShapedArray``
+            or ``jax.core.ShapedArray``
         """
-        return isinstance(x, ShapedArray)
+        return isinstance(x, (DShapedArray, ShapedArray))
+
+    @staticmethod
+    def is_dynamic_shape(shape):
+        """Determine whether a shape contains a tracer or not.
+
+        Args:
+            shape: The shape to be tested.
+
+        Returns:
+            bool: Whether the shape contains a tracer or not.
+        """
+        for s in shape:
+            if isinstance(s, Tracer):
+                return True
+
+        return False
 
     def __eq__(self, other):
         return self.xs == other.xs and self.ys == other.ys
@@ -123,8 +141,12 @@ def calculate_grad_shape(signature, indices) -> Signature:
                 grad_res_shape.append(axis)
             element_type = diff_arg_type.dtype
 
-            grad_res_type = (
-                ShapedArray(grad_res_shape, element_type) if grad_res_shape else diff_arg_type
-            )
+            grad_res_type = diff_arg_type
+            if grad_res_shape:
+                XShapedArray = (
+                    DShapedArray if Signature.is_dynamic_shape(grad_res_shape) else ShapedArray
+                )
+                grad_res_type = XShapedArray(grad_res_shape, element_type)
+
             grad_result_types.append(grad_res_type)
     return Signature(signature.get_inputs(), grad_result_types)

frontend/test/pytest/test_measurement_primitives.py

@@ -16,6 +16,7 @@
 """
 import jax
 
+from catalyst.jax_extras import DynamicJaxprTracer
 from catalyst.jax_primitives import (
     compbasis_p,
     counts_p,
@@ -40,6 +41,25 @@ def f():
     assert jaxpr.eqns[1].outvars[0].aval.shape == (5, 0)
 
 
+def test_sample_dynamic():
+    """Test that the sample primitive can be captured into jaxpr
+    when using tracers."""
+
+    def f(shots):
+        obs = compbasis_p.bind()
+        return sample_p.bind(obs, shots=shots, shape=(shots, 0))
+
+    jaxpr = jax.make_jaxpr(f)(5).jaxpr
+    shots_tracer, shape_value = jaxpr.eqns[1].params.values()
+
+    assert jaxpr.eqns[1].primitive == sample_p
+    assert isinstance(shots_tracer, DynamicJaxprTracer)
+    assert isinstance(shape_value[0], DynamicJaxprTracer)
+    assert shape_value[1] == 0
+    assert isinstance(jaxpr.eqns[1].outvars[0].aval.shape[0], jax.core.Var)
+    assert jaxpr.eqns[1].outvars[0].aval.shape[1] == 0
+
+
 def test_counts():
     """Test that the counts primitive can be captured by jaxpr."""
 
@@ -54,6 +74,24 @@ def f():
     assert jaxpr.eqns[1].outvars[1].aval.shape == (1,)
 
 
+def test_counts_dynamic():
+    """Test that the counts primitive can be captured by jaxpr
+    when using tracers."""
+
+    def f(dim_0):
+        obs = compbasis_p.bind()
+        return counts_p.bind(obs, shots=5, shape=(dim_0,))
+
+    jaxpr = jax.make_jaxpr(f)(1)
+    shots, shape_value = jaxpr.eqns[1].params.values()
+
+    assert jaxpr.eqns[1].primitive == counts_p
+    assert isinstance(shape_value[0], DynamicJaxprTracer)
+    assert shots == 5
+    assert isinstance(jaxpr.eqns[1].outvars[0].aval.shape[0], jax.core.Var)
+    assert isinstance(jaxpr.eqns[1].outvars[1].aval.shape[0], jax.core.Var)
+
+
 def test_expval():
     """Test that the expval primitive can be captured by jaxpr."""
 
@@ -81,7 +119,7 @@ def f():
 
 
 def test_probs():
-    """Test that the var primitive can be captured by jaxpr."""
+    """Test that the probs primitive can be captured by jaxpr."""
 
     def f():
         obs = compbasis_p.bind()
@@ -93,6 +131,23 @@ def f():
     assert jaxpr.eqns[1].outvars[0].aval.shape == (1,)
 
 
+def test_probs_dynamic():
+    """Test that the probs primitive can be captured by jaxpr
+    when using tracers."""
+
+    def f(dim_0):
+        obs = compbasis_p.bind()
+        return probs_p.bind(obs, shots=5, shape=(dim_0,))
+
+    jaxpr = jax.make_jaxpr(f)(1)
+    shots, shape_value = jaxpr.eqns[1].params.values()
+
+    assert jaxpr.eqns[1].primitive == probs_p
+    assert isinstance(shape_value[0], DynamicJaxprTracer)
+    assert shots == 5
+    assert isinstance(jaxpr.eqns[1].outvars[0].aval.shape[0], jax.core.Var)
+
+
 def test_state():
     """Test that the state primitive can be captured by jaxpr."""
 
@@ -104,3 +159,20 @@ def f():
     assert jaxpr.eqns[1].primitive == state_p
     assert jaxpr.eqns[1].params == {"shape": (1,), "shots": 5}
     assert jaxpr.eqns[1].outvars[0].aval.shape == (1,)
+
+
+def test_state_dynamic():
+    """Test that the state primitive can be captured by jaxpr
+    when using tracers."""
+
+    def f(dim_0):
+        obs = compbasis_p.bind()
+        return state_p.bind(obs, shots=5, shape=(dim_0,))
+
+    jaxpr = jax.make_jaxpr(f)(1)
+    shots, shape_value = jaxpr.eqns[1].params.values()
+
+    assert jaxpr.eqns[1].primitive == state_p
+    assert isinstance(shape_value[0], DynamicJaxprTracer)
+    assert shots == 5
+    assert isinstance(jaxpr.eqns[1].outvars[0].aval.shape[0], jax.core.Var)