Repository: PennyLaneAI/pennylane
Pull request #3534 (2022-12-14 14:32:52)

Fix jax
Branch out from this [PR's commit](https://github.com/PennyLaneAI/pennylane/pull/3531) to merge the changes into master.

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3534](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (4920f09) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/65440be09617b0ed967c04f0d332b5655ee376cd?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (65440be) will **not change** coverage.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3534   +/-   ##
=======================================
  Coverage   99.87%   99.87%           
=======================================
  Files         311      311           
  Lines       27899    27899           
=======================================
  Hits        27863    27863           
  Misses         36       36           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/gradients/finite\_difference.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9maW5pdGVfZGlmZmVyZW5jZS5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/gradients/gradient\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9ncmFkaWVudF90cmFuc2Zvcm0ucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/gradients/hessian\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9oZXNzaWFuX3RyYW5zZm9ybS5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/gradients/jvp.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9qdnAucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/gradients/parameter\_shift.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnQucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `100.00% <ø> (ø)` | |
| [pennylane/interfaces/jax\_jit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4X2ppdC5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `99.11% <ø> (ø)` | |
| [pennylane/qnode.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fub2RlLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/return\_types.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3JldHVybl90eXBlcy5weQ==) | `100.00% <ø> (ø)` | |
| ... and [2 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3534/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/doc/development/adding_operators.rst b/doc/development/adding_operators.rst
index 7326fb40c..d8d1de5a9 100644
--- a/doc/development/adding_operators.rst
+++ b/doc/development/adding_operators.rst
@@ -42,12 +42,12 @@ The basic components of operators are the following:
 
 #. **Trainable parameters** (:attr:`.Operator.parameters`) that the map depends on, such as a rotation angle,
    which can be fed to the operator as tensor-like objects. For example, since we used jax arrays to
-   specify the three rotation angles of ``op``, the parameters are jax ``DeviceArrays``.
+   specify the three rotation angles of ``op``, the parameters are jax ``Arrays``.
 
    >>> op.parameters
-   [DeviceArray(0.1, dtype=float32, weak_type=True),
-    DeviceArray(0.2, dtype=float32, weak_type=True),
-    DeviceArray(0.3, dtype=float32, weak_type=True)]
+   [Array(0.1, dtype=float32, weak_type=True),
+    Array(0.2, dtype=float32, weak_type=True),
+    Array(0.3, dtype=float32, weak_type=True)]
 
 #. **Non-trainable hyperparameters** (:attr:`.Operator.hyperparameters`) that influence the action of the operator.
    Not every operator has hyperparameters.
diff --git a/doc/development/guide/architecture.rst b/doc/development/guide/architecture.rst
index 08af71777..d6997a270 100644
--- a/doc/development/guide/architecture.rst
+++ b/doc/development/guide/architecture.rst
@@ -67,9 +67,9 @@ These four defining properties are accessible for all :class:`~.Operator` instan
 >>> op.name
 Rot
 >>> op.parameters
-[DeviceArray(0.1, dtype=float32, weak_type=True),
- DeviceArray(0.2, dtype=float32, weak_type=True),
- DeviceArray(0.3, dtype=float32, weak_type=True)]
+[Array(0.1, dtype=float32, weak_type=True),
+ Array(0.2, dtype=float32, weak_type=True),
+ Array(0.3, dtype=float32, weak_type=True)]
 >>> op.hyperparameters
 {}
 >>> op.wires
@@ -147,9 +147,9 @@ to the tape's ``measurement`` property.
 ...	    qfunc(params)
 
 >>> tape.operations
-[RX(DeviceArray(0.5, dtype=float32), wires=['b']),
+[RX(Array(0.5, dtype=float32), wires=['b']),
  CNOT(wires=['a', 'b']),
- RY(DeviceArray(0.2, dtype=float32), wires=['a'])]
+ RY(Array(0.2, dtype=float32), wires=['a'])]
 
 >>> tape.measurements
 [expval(PauliZ(wires=['b']))]
diff --git a/doc/introduction/interfaces/jax.rst b/doc/introduction/interfaces/jax.rst
index 340a2a6a7..9633a0206 100644
--- a/doc/introduction/interfaces/jax.rst
+++ b/doc/introduction/interfaces/jax.rst
@@ -67,13 +67,13 @@ a JAX-capable QNode in PennyLane. Simply specify the ``interface='jax'`` keyword
         qml.PhaseShift(theta, wires=0)
         return qml.expval(qml.PauliZ(0)), qml.expval(qml.Hadamard(1))
 
-The QNode ``circuit1()`` is now a JAX-capable QNode, accepting ``jax.DeviceArray`` objects
-as input, and returning ``jax.DeviceArray`` objects. It can now be used like any other JAX function:
+The QNode ``circuit1()`` is now a JAX-capable QNode, accepting ``jax.Array`` objects
+as input, and returning ``jax.Array`` objects. It can now be used like any other JAX function:
 
 >>> phi = jnp.array([0.5, 0.1])
 >>> theta = jnp.array(0.2)
 >>> circuit1(phi, theta)
-DeviceArray([0.8776, 0.6880], dtype=float64)
+Array([0.8776, 0.6880], dtype=float64)
 
 Quantum gradients using JAX
 ---------------------------
@@ -103,9 +103,9 @@ For example:
 This has output:
 
 >>> phi_grad
-DeviceArray([-0.47942555,  0.        ], dtype=float32)
+Array([-0.47942555,  0.        ], dtype=float32)
 >>> theta_grad
-DeviceArray(-3.4332792e-10, dtype=float32)
+Array(-3.4332792e-10, dtype=float32)
 
 
 .. _jax_jit:
@@ -234,7 +234,7 @@ used to optimize a QNode that is transformed by ``jax.jit``:
     optimized_params = res.params
 
 >>> optimized_params
-DeviceArray(3.1415861, dtype=float64, weak_type=True)
+Array(3.1415861, dtype=float64, weak_type=True)
 
 Alternatively, optimizers from ``Optax`` may also be used to optimize the same
 QNode:
@@ -267,4 +267,4 @@ QNode:
         params = optax.apply_updates(params, updates)
 
 >>> params
-DeviceArray(3.14159111, dtype=float64)
+Array(3.14159111, dtype=float64)
diff --git a/doc/releases/changelog-0.18.0.md b/doc/releases/changelog-0.18.0.md
index a95b7253a..f528ceede 100644
--- a/doc/releases/changelog-0.18.0.md
+++ b/doc/releases/changelog-0.18.0.md
@@ -478,7 +478,7 @@
   def group(coeffs, select=None):
     _, grouped_coeffs = qml.grouping.group_observables(obs, coeffs)
     # in this example, grouped_coeffs is a list of two jax tensors
-    # [DeviceArray([1., 2.], dtype=float32), DeviceArray([3.], dtype=float32)]
+    # [Array([1., 2.], dtype=float32), Array([3.], dtype=float32)]
     return grouped_coeffs[select]
 
   jac_fn = jax.jacobian(group)
diff --git a/doc/releases/changelog-0.21.0.md b/doc/releases/changelog-0.21.0.md
index 9defcc03e..a380e6dfd 100644
--- a/doc/releases/changelog-0.21.0.md
+++ b/doc/releases/changelog-0.21.0.md
@@ -218,10 +218,10 @@
   The QNode can be evaluated and its jacobian can be computed:
   ```pycon
   >>> circuit(x, y)
-  DeviceArray([0.8397495 , 0.16025047], dtype=float32)
+  Array([0.8397495 , 0.16025047], dtype=float32)
   >>> jax.jacobian(circuit, argnums=[0, 1])(x, y)
-  (DeviceArray([-0.2050439,  0.2050439], dtype=float32, weak_type=True),
-   DeviceArray([ 0.26043, -0.26043], dtype=float32, weak_type=True))
+  (Array([-0.2050439,  0.2050439], dtype=float32, weak_type=True),
+   Array([ 0.26043, -0.26043], dtype=float32, weak_type=True))
   ```
   Note that `jax.jit` is not yet supported for vector-valued QNodes.
 
diff --git a/doc/releases/changelog-0.22.0.md b/doc/releases/changelog-0.22.0.md
index a46f1187b..18fc233a1 100644
--- a/doc/releases/changelog-0.22.0.md
+++ b/doc/releases/changelog-0.22.0.md
@@ -607,7 +607,7 @@
 * The `qml.QubitUnitary` operation now supports just-in-time compilation using JAX.
   [(#2249)](https://github.com/PennyLaneAI/pennylane/pull/2249)
 
-* Fixes a bug in the JAX interface where `DeviceArray` objects
+* Fixes a bug in the JAX interface where `Array` objects
   were not being converted to NumPy arrays before executing an
   external device.
   [(#2255)](https://github.com/PennyLaneAI/pennylane/pull/2255)
diff --git a/doc/releases/changelog-0.23.0.md b/doc/releases/changelog-0.23.0.md
index 42fca99e7..482037a72 100644
--- a/doc/releases/changelog-0.23.0.md
+++ b/doc/releases/changelog-0.23.0.md
@@ -602,7 +602,7 @@
 
 * Fixed a bug where computing statistics for a relatively few number of shots
   (e.g., `shots=10`), an error arose due to indexing into an array using a
-  `DeviceArray`.
+  `Array`.
   [(#2427)](https://github.com/PennyLaneAI/pennylane/pull/2427)
 
 * PennyLane Lightning version in Docker container is pulled from latest wheel-builds.
diff --git a/doc/releases/changelog-0.24.0.md b/doc/releases/changelog-0.24.0.md
index 0d1945cf2..9f9ffc573 100644
--- a/doc/releases/changelog-0.24.0.md
+++ b/doc/releases/changelog-0.24.0.md
@@ -304,7 +304,7 @@
   ```
   ```pycon
   >>> circuit(x, y)
-  DeviceArray([0.8397495 , 0.16025047], dtype=float32)
+  Array([0.8397495 , 0.16025047], dtype=float32)
   ```
 
   Note that computing the jacobian of vector-valued QNode is not supported with JAX
@@ -322,7 +322,7 @@
   
   ```pycon
   >>> jax.grad(cost, argnums=[0])(x, y)
-  (DeviceArray(-0.2050439, dtype=float32),)
+  (Array(-0.2050439, dtype=float32),)
   ```
 
 <h4>More drawing styles 🎨</h4>
diff --git a/doc/releases/changelog-0.27.0.md b/doc/releases/changelog-0.27.0.md
index 4443c06fe..fc2a5a4a4 100644
--- a/doc/releases/changelog-0.27.0.md
+++ b/doc/releases/changelog-0.27.0.md
@@ -227,7 +227,7 @@
   Result value: 1.00; Result type: <class 'pennylane.numpy.tensor.tensor'>
   Result value: 1.00; Result type: <class 'torch.Tensor'>
   Result value: 1.00; Result type: <class 'tensorflow.python.framework.ops.EagerTensor'>
-  Result value: 1.00; Result type: <class 'jaxlib.xla_extension.DeviceArray'>
+  Result value: 1.00; Result type: <class 'jaxlib.xla_extension.Array'>
   ```
 
 <h4>Upgraded JAX-JIT gradient support 🏎</h4>
@@ -258,8 +258,8 @@
 
   ```pycon
   >>> jax.jacobian(circuit, argnums=[0, 1])(x, y)
-  (DeviceArray([-0.84147098,  0.35017549], dtype=float64, weak_type=True),
-   DeviceArray([ 4.47445479e-18, -4.91295496e-01], dtype=float64, weak_type=True))
+  (Array([-0.84147098,  0.35017549], dtype=float64, weak_type=True),
+   Array([ 4.47445479e-18, -4.91295496e-01], dtype=float64, weak_type=True))
   ```
 
   Note that this change depends on `jax.pure_callback`, which requires `jax>=0.3.17`.
@@ -364,15 +364,15 @@
 
   ```pycon
   >>> jax.hessian(circuit)(params)
-  ((DeviceArray([[ 0.,  0.],
+  ((Array([[ 0.,  0.],
                 [ 2., -3.]], dtype=float32),
-  DeviceArray([[[-0.5,  0. ],
+  Array([[[-0.5,  0. ],
                 [ 0. ,  0. ]],
               [[ 0.5,  0. ],
                 [ 0. ,  0. ]]], dtype=float32)),
-  (DeviceArray([[ 0.07677898,  0.0563341 ],
+  (Array([[ 0.07677898,  0.0563341 ],
                 [ 0.07238522, -1.830669  ]], dtype=float32),
-  DeviceArray([[[-4.9707499e-01,  2.9999996e-04],
+  Array([[[-4.9707499e-01,  2.9999996e-04],
                 [-6.2500127e-04,  1.2500001e-04]],
                 [[ 4.9707499e-01, -2.9999996e-04],
                 [ 6.2500127e-04, -1.2500001e-04]]], dtype=float32)))
diff --git a/doc/releases/changelog-0.28.0.md b/doc/releases/changelog-0.28.0.md
index c1bec7e50..f48f357a7 100644
--- a/doc/releases/changelog-0.28.0.md
+++ b/doc/releases/changelog-0.28.0.md
@@ -596,10 +596,10 @@
 
   ```pycon
   >>> jax.jacobian(circuit, argnums=[0, 1])(a, b)
-  ((DeviceArray(0.35017549, dtype=float64, weak_type=True),
-  DeviceArray(-0.4912955, dtype=float64, weak_type=True)),
-  (DeviceArray(5.55111512e-17, dtype=float64, weak_type=True),
-  DeviceArray(0., dtype=float64, weak_type=True)))
+  ((Array(0.35017549, dtype=float64, weak_type=True),
+  Array(-0.4912955, dtype=float64, weak_type=True)),
+  (Array(5.55111512e-17, dtype=float64, weak_type=True),
+  Array(0., dtype=float64, weak_type=True)))
   ```
 
 * Updated `qml.transforms.split_non_commuting` to support the new return types.
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
index 96426af65..90ed1c20c 100644
--- a/pennylane/gradients/finite_difference.py
+++ b/pennylane/gradients/finite_difference.py
@@ -294,8 +294,8 @@ def _finite_diff_new(
     ...     return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))
     >>> params = jax.numpy.array([0.1, 0.2, 0.3])
     >>> jax.jacobian(circuit)(params)
-    (DeviceArray([-0.38751727, -0.18884793, -0.3835571 ], dtype=float32),
-    DeviceArray([0.6991687 , 0.34072432, 0.6920237 ], dtype=float32))
+    (Array([-0.38751727, -0.18884793, -0.3835571 ], dtype=float32),
+    Array([0.6991687 , 0.34072432, 0.6920237 ], dtype=float32))
 
 
     .. details::
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
index caf32a6c4..60d48ff94 100644
--- a/pennylane/gradients/gradient_transform.py
+++ b/pennylane/gradients/gradient_transform.py
@@ -230,7 +230,7 @@ class gradient_transform(qml.batch_transform):
     .. note::
 
         The input tape might have parameters of various types, including
-        NumPy arrays, JAX DeviceArrays, and TensorFlow and PyTorch tensors.
+        NumPy arrays, JAX Arrays, and TensorFlow and PyTorch tensors.
 
         If the gradient transform is written in a autodiff-compatible manner, either by
         using a framework such as Autograd or TensorFlow, or by using ``qml.math`` for
diff --git a/pennylane/gradients/hessian_transform.py b/pennylane/gradients/hessian_transform.py
index d7cc6a640..231f77821 100644
--- a/pennylane/gradients/hessian_transform.py
+++ b/pennylane/gradients/hessian_transform.py
@@ -134,7 +134,7 @@ class hessian_transform(qml.batch_transform):
     .. note::
 
         The input tape might have parameters of various types, including NumPy arrays,
-        JAX DeviceArrays, and TensorFlow and PyTorch tensors.
+        JAX Arrays, and TensorFlow and PyTorch tensors.
 
         If the Hessian transform is written in a autodiff-compatible manner, either by
         using a framework such as Autograd or TensorFlow, or by using ``qml.math`` for
diff --git a/pennylane/gradients/jvp.py b/pennylane/gradients/jvp.py
index 90074a6cf..8777b0ce8 100644
--- a/pennylane/gradients/jvp.py
+++ b/pennylane/gradients/jvp.py
@@ -225,7 +225,7 @@ def jvp(tape, tangent, gradient_fn, shots=None, gradient_kwargs=None):
     >>> dev = qml.device("default.qubit", wires=2)
     >>> jvp = fn(dev.batch_execute(jvp_tapes))
     >>> jvp
-    (DeviceArray(-0.62073976, dtype=float32), DeviceArray([-0.3259707 ,  0.32597077], dtype=float32))
+    (Array(-0.62073976, dtype=float32), Array([-0.3259707 ,  0.32597077], dtype=float32))
     """
     gradient_kwargs = gradient_kwargs or {}
     num_params = len(tape.trainable_params)
@@ -357,7 +357,7 @@ def batch_jvp(tapes, tangents, gradient_fn, shots=None, reduction="append", grad
     >>> dev = qml.device("default.qubit", wires=2)
     >>> jvps = fn(dev.batch_execute(jvp_tapes))
     >>> jvps
-    [(DeviceArray(-0.62073976, dtype=float32), DeviceArray([-0.3259707 ,  0.32597077], dtype=float32)), DeviceArray(-0.6900841, dtype=float32)]
+    [(Array(-0.62073976, dtype=float32), Array([-0.3259707 ,  0.32597077], dtype=float32)), Array(-0.6900841, dtype=float32)]
 
     We have two JVPs; one per tape. Each one corresponds to the shape of the output of their respective tape.
     """
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index 371958dce..fe27b8ed9 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -1340,7 +1340,7 @@ def _param_shift_new(
     ...     return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))
     >>> params = jax.numpy.array([0.1, 0.2, 0.3])
     >>> jax.jacobian(circuit)(params)
-    (DeviceArray([-0.38751727, -0.18884793, -0.3835571 ], dtype=float32), DeviceArray([0.6991687 , 0.34072432, 0.6920237 ], dtype=float32))
+    (Array([-0.38751727, -0.18884793, -0.3835571 ], dtype=float32), Array([0.6991687 , 0.34072432, 0.6920237 ], dtype=float32))
 
     .. note::
 
diff --git a/pennylane/interfaces/jax.py b/pennylane/interfaces/jax.py
index 3c84f9f9a..739bab620 100644
--- a/pennylane/interfaces/jax.py
+++ b/pennylane/interfaces/jax.py
@@ -246,11 +246,11 @@ def _execute(
             # Unstack partial results into ndim=0 arrays to allow
             # differentiability with JAX
             # E.g.,
-            # [DeviceArray([-0.9553365], dtype=float32), DeviceArray([0., 0.],
+            # [Array([-0.9553365], dtype=float32), Array([0., 0.],
             # dtype=float32)]
             # is mapped to
-            # [[DeviceArray(-0.9553365, dtype=float32)], [DeviceArray(0.,
-            # dtype=float32), DeviceArray(0., dtype=float32)]].
+            # [[Array(-0.9553365, dtype=float32)], [Array(0.,
+            # dtype=float32), Array(0., dtype=float32)]].
             need_unstacking = any(r.ndim != 0 for r in res)
             if need_unstacking:
                 res = [qml.math.unstack(x) for x in res]
@@ -281,13 +281,13 @@ def _raise_vector_valued_fwd(tapes):
     Example to the latter:
 
     1. Output when using jax.jacobian:
-    DeviceArray([[-0.09983342,  0.01983384],\n
+    Array([[-0.09983342,  0.01983384],\n
                  [-0.09983342, 0.01983384]], dtype=float64),
-    DeviceArray([[ 0.        , -0.97517033],\n
+    Array([[ 0.        , -0.97517033],\n
                  [ 0.        , -0.97517033]], dtype=float64)),
 
     2. Expected output:
-    DeviceArray([[-0.09983342, 0.01983384],\n
+    Array([[-0.09983342, 0.01983384],\n
                 [ 0.        , -0.97517033]]
 
     The output produced by this function matches 1.
@@ -342,10 +342,10 @@ def _execute_fwd(
         # Adjust the structure of how the jacobian is returned to match the
         # non-forward mode cases
         # E.g.,
-        # [DeviceArray([[ 0.06695931,  0.01383095, -0.46500877]], dtype=float32)]
+        # [Array([[ 0.06695931,  0.01383095, -0.46500877]], dtype=float32)]
         # is mapped to
-        # [[DeviceArray(0.06695931, dtype=float32), DeviceArray(0.01383095,
-        # dtype=float32), DeviceArray(-0.46500877, dtype=float32)]]
+        # [[Array(0.06695931, dtype=float32), Array(0.01383095,
+        # dtype=float32), Array(-0.46500877, dtype=float32)]]
         res_jacs = []
         for j in jacs:
             this_j = []
diff --git a/pennylane/interfaces/jax_jit.py b/pennylane/interfaces/jax_jit.py
index 0cb4292e4..583ecc70c 100644
--- a/pennylane/interfaces/jax_jit.py
+++ b/pennylane/interfaces/jax_jit.py
@@ -240,11 +240,11 @@ def _execute(
             # Unwrap partial results into ndim=0 arrays to allow
             # differentiability with JAX
             # E.g.,
-            # [DeviceArray([-0.9553365], dtype=float32), DeviceArray([0., 0.],
+            # [Array([-0.9553365], dtype=float32), Array([0., 0.],
             # dtype=float32)]
             # is mapped to
-            # [[DeviceArray(-0.9553365, dtype=float32)], [DeviceArray(0.,
-            # dtype=float32), DeviceArray(0., dtype=float32)]].
+            # [[Array(-0.9553365, dtype=float32)], [Array(0.,
+            # dtype=float32), Array(0., dtype=float32)]].
             need_unstacking = any(r.ndim != 0 for r in res)
             if need_unstacking:
                 res = [qml.math.unstack(x) for x in res]
@@ -328,10 +328,10 @@ def _execute_with_fwd(
         # Adjust the structure of how the jacobian is returned to match the
         # non-forward mode cases
         # E.g.,
-        # [DeviceArray([[ 0.06695931,  0.01383095, -0.46500877]], dtype=float32)]
+        # [Array([[ 0.06695931,  0.01383095, -0.46500877]], dtype=float32)]
         # is mapped to
-        # [[DeviceArray(0.06695931, dtype=float32), DeviceArray(0.01383095,
-        # dtype=float32), DeviceArray(-0.46500877, dtype=float32)]]
+        # [[Array(0.06695931, dtype=float32), Array(0.01383095,
+        # dtype=float32), Array(-0.46500877, dtype=float32)]]
         res_jacs = []
         for j in jacs:
             this_j = []
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
index 15927bca7..2895521ed 100644
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -404,11 +404,11 @@ def get_trainable_indices(values, like=None):
         import jax
 
         if not any(isinstance(v, jax.core.Tracer) for v in values):
-            # No JAX tracing is occuring; treat all `DeviceArray` objects as trainable.
+            # No JAX tracing is occuring; treat all `Array` objects as trainable.
 
             # pylint: disable=function-redefined,unused-argument
             def trainable(p, **kwargs):
-                return isinstance(p, jax.numpy.DeviceArray)
+                return isinstance(p, jax.Array)
 
         else:
             # JAX tracing is occuring; use the default behaviour (only traced arrays
diff --git a/pennylane/math/utils.py b/pennylane/math/utils.py
index 4c8af9fc1..ff9a53512 100644
--- a/pennylane/math/utils.py
+++ b/pennylane/math/utils.py
@@ -322,7 +322,7 @@ def is_abstract(tensor, like=None):
     >>> function(x)
     Value: [0.5, 0.1]
     Abstract: False
-    DeviceArray(0.26, dtype=float32)
+    Array(0.26, dtype=float32)
 
     However, if we use the ``@jax.jit`` decorator, the tensor will now be abstract:
 
@@ -330,7 +330,7 @@ def is_abstract(tensor, like=None):
     >>> jax.jit(function)(x)
     Value: Traced<ShapedArray(float32[2])>with<DynamicJaxprTrace(level=0/1)>
     Abstract: True
-    DeviceArray(0.26, dtype=float32)
+    Array(0.26, dtype=float32)
 
     Note that JAX uses an abstract *shaped* array, so although we won't be able to
     include conditionals within our function that depend on the value of the tensor,
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 8fb5afc71..79d50ed25 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -60,7 +60,7 @@ class QNode:
 
             * ``"jax"``: Allows JAX to backpropagate
               through the QNode. The QNode accepts and returns
-              JAX ``DeviceArray`` objects.
+              JAX ``Array`` objects.
 
             * ``None``: The QNode accepts default Python types
               (floats, ints, lists, tuples, dicts) as well as NumPy array arguments,
diff --git a/pennylane/return_types.py b/pennylane/return_types.py
index 1d7218248..3f44d7f9c 100644
--- a/pennylane/return_types.py
+++ b/pennylane/return_types.py
@@ -204,10 +204,10 @@ def enable_return():
                 return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
 
         >>> jax.hessian(circuit, argnums=[0, 1])(par_0, par_1)
-        ((DeviceArray(-0.19767681, dtype=float32, weak_type=True),
-          DeviceArray(-0.09784342, dtype=float32, weak_type=True)),
-         (DeviceArray(-0.09784339, dtype=float32, weak_type=True),
-          DeviceArray(-0.19767687, dtype=float32, weak_type=True)))
+        ((Array(-0.19767681, dtype=float32, weak_type=True),
+          Array(-0.09784342, dtype=float32, weak_type=True)),
+         (Array(-0.09784339, dtype=float32, weak_type=True),
+          Array(-0.19767687, dtype=float32, weak_type=True)))
 
         The new return types system also unlocks the use of ``probs`` mixed with different measurements with JAX:
 
@@ -231,8 +231,8 @@ def enable_return():
             x = jax.numpy.array([0.1, 0.2, 0.3])
 
         >>> jax.jacobian(circuit)(x)
-        (DeviceArray([-9.9833414e-02, -7.4505806e-09,  6.9285655e-10], dtype=float32),
-         DeviceArray([[-4.9419206e-02, -9.9086545e-02,  3.4938008e-09],
+        (Array([-9.9833414e-02, -7.4505806e-09,  6.9285655e-10], dtype=float32),
+         Array([[-4.9419206e-02, -9.9086545e-02,  3.4938008e-09],
                       [-4.9750542e-04,  9.9086538e-02,  1.2768372e-10],
                       [ 4.9750548e-04,  2.4812977e-04,  4.8371929e-13],
                       [ 4.9419202e-02, -2.4812980e-04,  2.6696912e-11]],            dtype=float32))
@@ -264,15 +264,15 @@ def enable_return():
                 return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         >>> jax.hessian(circuit)(params)
-        ((DeviceArray([[ 0.,  0.],
+        ((Array([[ 0.,  0.],
                         [ 2., -3.]], dtype=float32),
-          DeviceArray([[[-0.5,  0. ],
+          Array([[[-0.5,  0. ],
                        [ 0. ,  0. ]],
                       [[ 0.5,  0. ],
                        [ 0. ,  0. ]]], dtype=float32)),
-         (DeviceArray([[ 0.07677898,  0.0563341 ],
+         (Array([[ 0.07677898,  0.0563341 ],
                        [ 0.07238522, -1.830669  ]], dtype=float32),
-          DeviceArray([[[-4.9707499e-01,  2.9999996e-04],
+          Array([[[-4.9707499e-01,  2.9999996e-04],
                         [-6.2500127e-04,  1.2500001e-04]],
                        [[ 4.9707499e-01, -2.9999996e-04],
                         [ 6.2500127e-04, -1.2500001e-04]]], dtype=float32)))
diff --git a/pennylane/transforms/decompositions/two_qubit_unitary.py b/pennylane/transforms/decompositions/two_qubit_unitary.py
index 01f8b3aa8..5e782399e 100644
--- a/pennylane/transforms/decompositions/two_qubit_unitary.py
+++ b/pennylane/transforms/decompositions/two_qubit_unitary.py
@@ -39,7 +39,7 @@ from .single_qubit_unitary import zyz_decomposition
 # - In Tensorflow, it sometimes works in limited cases (0, sometimes 1 CNOT), but
 #   for others it fails without output making it hard to pinpoint the cause.
 # - In JAX, we receive the TypeError:
-#       Can't differentiate w.r.t. type <class 'jaxlib.xla_extension.DeviceArray'>
+#       Can't differentiate w.r.t. type <class 'jaxlib.xla_extension.Array'>
 #
 ###################################################################################
 


Test patch:
diff --git a/pennylane/devices/tests/test_measurements.py b/pennylane/devices/tests/test_measurements.py
index 726e22545..f06caa39e 100755
--- a/pennylane/devices/tests/test_measurements.py
+++ b/pennylane/devices/tests/test_measurements.py
@@ -1501,14 +1501,15 @@ class TestSampleMeasurement:
             """Dummy sampled measurement."""
 
             def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
-                return qml.math.sum(samples[..., self.wires])
+                return 1
 
         @qml.qnode(dev)
         def circuit():
             qml.PauliX(0)
             return MyMeasurement(wires=[0]), MyMeasurement(wires=[1])
 
-        assert qml.math.allequal(circuit(), [dev.shots, 0])
+        res = circuit()
+        assert qml.math.allequal(res, [1, 1])
 
     def test_sample_measurement_without_shots(self, device):
         """Test that executing a sampled measurement with ``shots=None`` raises an error."""
@@ -1521,7 +1522,7 @@ class TestSampleMeasurement:
             """Dummy sampled measurement."""
 
             def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
-                return qml.math.sum(samples[..., self.wires])
+                return 1
 
         @qml.qnode(dev)
         def circuit():
@@ -1561,14 +1562,18 @@ class TestStateMeasurement:
         """Test the execution of a custom state measurement."""
         dev = device(2)
 
+        if dev.shots is not None:
+            pytest.skip("Some plugins don't update state information when shots is not None.")
+
         class MyMeasurement(StateMeasurement):
             """Dummy state measurement."""
 
             def process_state(self, state, wire_order):
-                return qml.math.sum(state)
+                return 1
 
         @qml.qnode(dev)
         def circuit():
+            qml.PauliX(0)
             return MyMeasurement()
 
         assert circuit() == 1
@@ -1584,10 +1589,11 @@ class TestStateMeasurement:
             """Dummy state measurement."""
 
             def process_state(self, state, wire_order):
-                return qml.math.sum(state)
+                return 1
 
         @qml.qnode(dev)
         def circuit():
+            qml.PauliX(0)
             return MyMeasurement()
 
         with pytest.warns(
@@ -1622,13 +1628,14 @@ class TestCustomMeasurement:
             """Dummy measurement transform."""
 
             def process(self, qscript, device):
-                return {device.shots: len(qscript)}
+                return 1
 
         @qml.qnode(dev)
         def circuit():
+            qml.PauliX(0)
             return MyMeasurement()
 
-        assert circuit() == {dev.shots: len(circuit.tape)}
+        assert circuit() == 1
 
     def test_method_overriden_by_device(self, device):
         """Test that the device can override a measurement process."""
@@ -1642,6 +1649,7 @@ class TestCustomMeasurement:
 
         @qml.qnode(dev)
         def circuit():
+            qml.PauliX(0)
             return qml.classical_shadow(wires=0)
 
         circuit.device.measurement_map[ClassicalShadowMP] = "test_method"
diff --git a/tests/interfaces/test_jax_qnode.py b/tests/interfaces/test_jax_qnode.py
index 26a7a9c5e..4062a4b3d 100644
--- a/tests/interfaces/test_jax_qnode.py
+++ b/tests/interfaces/test_jax_qnode.py
@@ -75,7 +75,7 @@ class TestQNode:
 
         # gradients should work
         grad = jax.grad(circuit)(a)
-        assert isinstance(grad, jax.numpy.DeviceArray)
+        assert isinstance(grad, jax.Array)
         assert grad.shape == tuple()
 
     def test_changing_trainability(self, dev_name, diff_method, mode, interface, mocker, tol):
@@ -704,7 +704,7 @@ class TestQubitIntegration:
         res = circuit()
 
         assert res.shape == (2, 10)
-        assert isinstance(res, jax.numpy.DeviceArray)
+        assert isinstance(res, jax.Array)
 
     def test_chained_qnodes(self, dev_name, diff_method, mode, interface):
         """Test that the gradient of chained QNodes works without error"""
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
index 2a06e742d..5106c1c53 100644
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -1060,7 +1060,7 @@ class TestRequiresGrad:
 
     @pytest.mark.slow
     def test_jax(self):
-        """JAX DeviceArrays differentiability depends on the argnums argument"""
+        """JAX Arrays differentiability depends on the argnums argument"""
         res = None
 
         def cost_fn(t, s):
@@ -1156,7 +1156,7 @@ class TestInBackprop:
 
     @pytest.mark.slow
     def test_jax(self):
-        """The value of in_backprop for JAX DeviceArrays depends on the argnums argument"""
+        """The value of in_backprop for JAX Arrays depends on the argnums argument"""
         res = None
 
         def cost_fn(t, s):
@@ -1699,7 +1699,7 @@ class TestScatterElementAdd:
             return fn.scatter_element_add(weights[0], self.index, weights[1] ** 2)
 
         res = cost([x, y])
-        assert isinstance(res, jax.interpreters.xla.DeviceArray)
+        assert isinstance(res, jax.Array)
         assert fn.allclose(res, self.expected_val)
 
         grad = jax.grad(lambda weights: cost(weights)[self.index[0], self.index[1]])([x, y])
@@ -1716,7 +1716,7 @@ class TestScatterElementAdd:
             return fn.scatter_element_add(weight_0, self.index, weight_1**2)
 
         res = cost_multi(x, y)
-        assert isinstance(res, jax.interpreters.xla.DeviceArray)
+        assert isinstance(res, jax.Array)
         assert fn.allclose(res, self.expected_val)
 
         jac = jax.jacobian(lambda *weights: cost_multi(*weights), argnums=[0, 1])(x, y)
@@ -1812,7 +1812,7 @@ class TestScatterElementAddMultiValue:
             )
 
         res = cost([x, y])
-        assert isinstance(res, jax.interpreters.xla.DeviceArray)
+        assert isinstance(res, jax.Array)
         assert fn.allclose(res, self.expected_val)
 
         scalar_cost = (
diff --git a/tests/returntypes/jax/test_jax_jit_qnode_new.py b/tests/returntypes/jax/test_jax_jit_qnode_new.py
index e8b6db320..d21f1bbf2 100644
--- a/tests/returntypes/jax/test_jax_jit_qnode_new.py
+++ b/tests/returntypes/jax/test_jax_jit_qnode_new.py
@@ -64,7 +64,7 @@ class TestQNode:
 
         # gradients should work
         grad = jax.jit(jax.grad(circuit))(a)
-        assert isinstance(grad, jax.numpy.DeviceArray)
+        assert isinstance(grad, jax.Array)
         assert grad.shape == ()
 
     def test_changing_trainability(self, dev_name, diff_method, mode, interface, mocker, tol):
@@ -780,9 +780,9 @@ class TestQubitIntegration:
 
         assert isinstance(res, tuple)
 
-        assert isinstance(res[0], jax.numpy.DeviceArray)
+        assert isinstance(res[0], jax.Array)
         assert res[0].shape == (10,)
-        assert isinstance(res[1], jax.numpy.DeviceArray)
+        assert isinstance(res[1], jax.Array)
         assert res[1].shape == (10,)
 
     def test_counts(self, dev_name, diff_method, mode, interface):
diff --git a/tests/returntypes/jax/test_jax_qnode_new.py b/tests/returntypes/jax/test_jax_qnode_new.py
index f4bdaee81..cfb4cab6c 100644
--- a/tests/returntypes/jax/test_jax_qnode_new.py
+++ b/tests/returntypes/jax/test_jax_qnode_new.py
@@ -64,7 +64,7 @@ class TestQNode:
 
         # gradients should work
         grad = jax.grad(circuit)(a)
-        assert isinstance(grad, jax.numpy.DeviceArray)
+        assert isinstance(grad, jax.Array)
         assert grad.shape == ()
 
     def test_changing_trainability(self, dev_name, diff_method, mode, interface, mocker, tol):
@@ -771,9 +771,9 @@ class TestQubitIntegration:
 
         assert isinstance(res, tuple)
 
-        assert isinstance(res[0], jax.numpy.DeviceArray)
+        assert isinstance(res[0], jax.Array)
         assert res[0].shape == (10,)
-        assert isinstance(res[1], jax.numpy.DeviceArray)
+        assert isinstance(res[1], jax.Array)
         assert res[1].shape == (10,)
 
     def test_counts(self, dev_name, diff_method, mode, interface):
diff --git a/tests/tape/test_unwrap.py b/tests/tape/test_unwrap.py
index 181a78613..961a6e66a 100644
--- a/tests/tape/test_unwrap.py
+++ b/tests/tape/test_unwrap.py
@@ -202,7 +202,6 @@ def test_unwrap_jax_backward():
     works as expected during a backwards pass"""
     import jax
     from jax import numpy as jnp
-    from jaxlib.xla_extension import DeviceArray
     from jax.interpreters.ad import JVPTracer
 
     p = [
@@ -231,7 +230,7 @@ def test_unwrap_jax_backward():
 
         # outside the context, the original parameters have been restored.
         params = tape.get_parameters(trainable_only=False)
-        assert all(isinstance(i, (DeviceArray, JVPTracer)) for i in params)
+        assert all(isinstance(i, (jax.Array, JVPTracer)) for i in params)
 
         return p[0][0] * p[1] ** 2 * jnp.sin(p[0][1]) * jnp.exp(-0.5 * p[0][2])