Repository: PennyLaneAI/pennylane

Issue:
Jitting qnode fails with ParametrizedEvolution
### Expected behavior

Running the function with and without `jax.jit` would produce the same results when passing a single float for time

### Actual behavior

It raises a `TracerArrayConversionError` with jitting

### Additional information

The error goes away if in line 276 of parametrized_evolution.py, the list of times is cast as a jnp.array, i.e if:

            self.t = qml.math.array([0.0, t] if qml.math.ndim(t) == 0 else t, dtype=float)

becomes

            self.t = qml.math.array(jnp.array([0.0, t]) if qml.math.ndim(t) == 0 else t, dtype=float)


Converting to `jax.numpy.array([t])` when passing to qml.evolve also runs, but only because it skips the logic in `qml.math.ndim(t) == 0` (its ndim is 1), leaving time as an array of length 1. This returns an incorrect result.

### Source code

```shell
H = qml.pulse.ParametrizedHamiltonian([2], [qml.PauliX(0)])
    
dev = qml.device("default.qubit", wires=1)

@jax.jit
@qml.qnode(dev, interface="jax")
def circuit(t):
    qml.evolve(H)([], t)
    return qml.expval(qml.PauliZ(0))

circuit(1)
```


### Tracebacks

```shell
---------------------------------------------------------------------------
TracerArrayConversionError                Traceback (most recent call last)
Cell In[21], line 11
      8     qml.evolve(H)([], t)
      9     return qml.expval(qml.PauliZ(0))
---> 11 circuit(1)

    [... skipping hidden 12 frame]

File ~/pennylane/pennylane/qnode.py:847, in QNode.__call__(self, *args, **kwargs)
    844         set_shots(self._original_device, override_shots)(self._update_gradient_fn)()
    846 # construct the tape
--> 847 self.construct(args, kwargs)
    849 cache = self.execute_kwargs.get("cache", False)
    850 using_custom_cache = (
    851     hasattr(cache, "__getitem__")
    852     and hasattr(cache, "__setitem__")
    853     and hasattr(cache, "__delitem__")
    854 )

File ~/pennylane/pennylane/qnode.py:751, in QNode.construct(self, args, kwargs)
    748 if old_interface == "auto":
    749     self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
--> 751 self._tape = make_qscript(self.func)(*args, **kwargs)
    752 self._qfunc_output = self.tape._qfunc_output
    754 params = self.tape.get_parameters(trainable_only=False)

File ~/pennylane/pennylane/tape/qscript.py:1371, in make_qscript.<locals>.wrapper(*args, **kwargs)
   1369 def wrapper(*args, **kwargs):
   1370     with AnnotatedQueue() as q:
-> 1371         result = fn(*args, **kwargs)
   1373     qscript = QuantumScript.from_queue(q)
   1374     qscript._qfunc_output = result

Cell In[21], line 8, in circuit(t)
      5 @jax.jit
      6 @qml.qnode(dev, interface="jax")
      7 def circuit(t):
----> 8     qml.evolve(H)([], t)
      9     return qml.expval(qml.PauliZ(0))

File ~/pennylane/pennylane/pulse/parametrized_evolution.py:293, in ParametrizedEvolution.__call__(self, params, t, **odeint_kwargs)
    290 if qml.QueuingManager.recording():
    291     qml.QueuingManager.remove(self)
--> 293 return ParametrizedEvolution(
    294     H=self.H, params=params, t=t, do_queue=True, id=self.id, **odeint_kwargs
    295 )

File ~/pennylane/pennylane/pulse/parametrized_evolution.py:276, in ParametrizedEvolution.__init__(self, H, params, t, do_queue, id, **odeint_kwargs)
    274     self.t = None
    275 else:
--> 276     self.t = qml.math.array([0.0, t] if qml.math.ndim(t) == 0 else t, dtype=float)
    277 params = [] if params is None else params
    278 super().__init__(*params, wires=H.wires, do_queue=do_queue, id=id)

File ~/pennylane/pennylane/math/multi_dispatch.py:39, in array(like, *args, **kwargs)
     30 def array(*args, like=None, **kwargs):
     31     """Creates an array or tensor object of the target framework.
     32 
     33     If the PyTorch interface is specified, this method preserves the Torch device used.
   (...)
     37         tensor_like: the tensor_like object of the framework
     38     """
---> 39     res = np.array(*args, like=like, **kwargs)
     40     if like is not None and get_interface(like) == "torch":
     41         res = res.to(device=like.device)

File /opt/homebrew/Caskroom/miniforge/base/envs/pennylane/lib/python3.10/site-packages/autoray/autoray.py:79, in do(fn, like, *args, **kwargs)
     30 """Do function named ``fn`` on ``(*args, **kwargs)``, peforming single
     31 dispatch to retrieve ``fn`` based on whichever library defines the class of
     32 the ``args[0]``, or the ``like`` keyword argument if specified.
   (...)
     76     <tf.Tensor: id=91, shape=(3, 3), dtype=float32>
     77 """
     78 backend = choose_backend(fn, *args, like=like, **kwargs)
---> 79 return get_lib_fn(backend, fn)(*args, **kwargs)

File /opt/homebrew/Caskroom/miniforge/base/envs/pennylane/lib/python3.10/site-packages/jax/_src/core.py:557, in Tracer.__array__(self, *args, **kw)
    556 def __array__(self, *args, **kw):
--> 557   raise TracerArrayConversionError(self)

TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
The error occurred while tracing the function circuit at /var/folders/61/wr1fxnf95tg9k56bz1_7g29r0000gq/T/ipykernel_62432/3769722557.py:5 for jit. This concrete value was not available in Python because it depends on the value of the argument 't'.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError
```


### System information

```shell
pl dev
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:
The problem is the casting to an array, which does not work on jitted functions. We should be using `stack` instead. I'll have a fix up shortly.

Resolving patch:
diff --git a/pennylane/pulse/parametrized_evolution.py b/pennylane/pulse/parametrized_evolution.py
--- a/pennylane/pulse/parametrized_evolution.py
+++ b/pennylane/pulse/parametrized_evolution.py
@@ -385,7 +385,9 @@ def __init__(
         if t is None:
             self.t = None
         else:
-            self.t = qml.math.array([0.0, t] if qml.math.ndim(t) == 0 else t, dtype=float)
+            if isinstance(t, (list, tuple)):
+                t = qml.math.stack(t)
+            self.t = qml.math.cast(qml.math.stack([0.0, t]) if qml.math.ndim(t) == 0 else t, float)
         if complementary and not return_intermediate:
             warnings.warn(
                 "The keyword argument complementary does not have any effect if "


Test patch:
diff --git a/tests/gradients/core/test_pulse_gradient.py b/tests/gradients/core/test_pulse_gradient.py
--- a/tests/gradients/core/test_pulse_gradient.py
+++ b/tests/gradients/core/test_pulse_gradient.py
@@ -763,8 +763,9 @@ def circuit(params):
             for j, e in zip(jac, exp_jac):
                 assert qml.math.allclose(j[0], e, atol=tol, rtol=0.0)
 
+    @pytest.mark.xfail
     @pytest.mark.parametrize("num_split_times", [1, 2])
-    @pytest.mark.parametrize("time_interface", ["python", "numpy", "JAX"])
+    @pytest.mark.parametrize("time_interface", ["python", "numpy", "jax"])
     def test_simple_qnode_jit(self, num_split_times, time_interface):
         """Test that a simple qnode can be differentiated with stoch_pulse_grad."""
         import jax
@@ -772,20 +773,20 @@ def test_simple_qnode_jit(self, num_split_times, time_interface):
 
         jax.config.update("jax_enable_x64", True)
         dev = qml.device("default.qubit.jax", wires=1)
-        T = {"python": 0.2, "numpy": np.array(0.2), "JAX": jnp.array(0.2)}[time_interface]
+        T = {"python": 0.2, "numpy": np.array(0.2), "jax": jnp.array(0.2)}[time_interface]
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
         @qml.qnode(
             dev, interface="jax", diff_method=stoch_pulse_grad, num_split_times=num_split_times
         )
-        def circuit(params):
+        def circuit(params, T=None):
             qml.evolve(ham_single_q_const)(params, T)
             return qml.expval(qml.PauliZ(0))
 
         params = [jnp.array(0.4)]
         p = params[0] * T
         exp_grad = -2 * jnp.sin(2 * p) * T
-        jit_grad = jax.jit(jax.grad(circuit))(params)
+        jit_grad = jax.jit(jax.grad(circuit))(params, T=T)
         assert qml.math.isclose(jit_grad, exp_grad)
 
     @pytest.mark.slow
diff --git a/tests/pulse/test_parametrized_evolution.py b/tests/pulse/test_parametrized_evolution.py
--- a/tests/pulse/test_parametrized_evolution.py
+++ b/tests/pulse/test_parametrized_evolution.py
@@ -348,6 +348,33 @@ def test_return_intermediate_and_complementary(self, comp, len_t):
 class TestIntegration:
     """Integration tests for the ParametrizedEvolution class."""
 
+    @pytest.mark.parametrize("time", [0.3, 1, [0, 2], [0.4, 2], (3, 3.1)])
+    @pytest.mark.parametrize("time_interface", ["python", "numpy", "jax"])
+    @pytest.mark.parametrize("use_jit", [False, True])
+    def test_time_input_formats(self, time, time_interface, use_jit):
+        import jax
+        import jax.numpy as jnp
+
+        if time_interface == "jax":
+            time = jnp.array(time)
+        elif time_interface == "numpy":
+            time = np.array(time)
+        H = qml.pulse.ParametrizedHamiltonian([2], [qml.PauliX(0)])
+
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev, interface="jax")
+        def circuit(t):
+            qml.evolve(H)([], t)
+            return qml.expval(qml.PauliZ(0))
+
+        if use_jit:
+            circuit = jax.jit(circuit)
+
+        res = circuit(time)
+        duration = time if qml.math.ndim(time) == 0 else time[1] - time[0]
+        assert qml.math.isclose(res, qml.math.cos(4 * duration))
+
     # pylint: disable=unused-argument
     def test_time_independent_hamiltonian(self):
         """Test the execution of a time independent hamiltonian."""
diff --git a/tests/pulse/test_rydberg.py b/tests/pulse/test_rydberg.py
--- a/tests/pulse/test_rydberg.py
+++ b/tests/pulse/test_rydberg.py
@@ -19,9 +19,6 @@
 import numpy as np
 import pytest
 
-jax = pytest.importorskip("jax")
-import jax.numpy as jnp
-
 import pennylane as qml
 from pennylane.pulse import rydberg_interaction, rydberg_drive
 from pennylane.pulse.hardware_hamiltonian import (
@@ -417,8 +414,11 @@ def qnode_jit(params):
         assert isinstance(res, jax.Array)
         assert np.allclose(res, res_jit)
 
+    @pytest.mark.jax
     def test_pennylane_and_exact_solution_correspond(self):
         """Test that the results of PennyLane simulation match (within reason) the exact solution"""
+        import jax
+        import jax.numpy as jnp
 
         def exact(H, H_obj, t):
             psi0 = jnp.eye(2 ** len(H.wires))[0]