From 94f067a257d1e2a96812d30b576bc90da5381318 Mon Sep 17 00:00:00 2001
From: Cristian Emiliano Godinez Ramirez
 <57567043+EmilianoG-byte@users.noreply.github.com>
Date: Mon, 16 Sep 2024 15:08:41 +0200
Subject: [PATCH] Pennylane is compatible with numpy 2.0 (#6061)

**Context:**

We want to make Pennylane compatible with Numpy 2.0. After several
discussions, we decided to test NumPy 2.0 on the CI by default in every
PR (testing both Python versions would have been to slow). Some jobs
still downgrade automatically to Numpy 1.x, since some interfaces (such
as Tensorflow) still do not support NumPy 2.0.

**Description of the Change:**

We can distinguish the changes into 3 main categories:

*Changes to workflows*

- None in the final version

*Changes to requirements and setup files*

- Unpin the Numpy version in `setup.py` (now we also allow Numpy 2.0).
- Update `requirements-ci.txt` to include Scipy 1.13 (this adds support
for Numpy 2.0).
- Pin Numpy in `requirements-ci.txt` to 2.0.

*Changes to the source code*
- Change `np.NaN` to `np.nan`.
- Use legacy printing representation in tests, contrary to the new numpy
representation of scalars, e.g. np.float64(3.0) rather than just 3.0.
- Update probabilities warning to be case insensitive and check for a
partial match, since this warning was changed in Numpy 2.0.
- Check the datatype of np.exp from the Global phase only for Numpy 1.x,
since this gets promoted to complex128 in Numpy 2.x.
https://numpy.org/neps/nep-0050-scalar-promotion.html#schema-of-the-new-proposed-promotion-rules.

**Benefits:**

Make Pennylane compatible with Numpy 2.0.

**Possible Drawbacks:**

- We need to create a separate workflow to keep testing PennyLane with
NumPy 1.x, since we still want to maintain compatibility with previous
NumPy versions. This will be done in a separate PR.

- We are not testing Numpy 2.x for the interfaces that implicitly
require Numpy 1.x. These currently seem to be `tensorflow` and
`openfermionpyscf` (notice that `tensorflow` is required in some code
sections like qcut). In particular, `openfermionpyscf` causes an error:
```
AttributeError: np.string_ was removed in the NumPy 2.0 release. Use np.bytes_ instead.
```
in the qchem tests. The attribute `np.string_` is not used in the PL
source code, so it is a problem with the package itself.

[sc-61399]
[sc-66548]

---------

Co-authored-by: PietropaoloFrisoni <pietropaolo.frisoni@xanadu.ai>
Co-authored-by: Pietropaolo Frisoni <pietropfrisoni@gmail.com>
---
 .github/workflows/install_deps/action.yml     |  6 +--
 .gitignore                                    |  2 +
 doc/releases/changelog-dev.md                 |  5 +++
 pennylane/numpy/random.py                     |  5 ++-
 requirements-ci.txt                           |  2 +-
 setup.py                                      |  2 +-
 .../data/attributes/operator/test_operator.py | 11 +++--
 tests/data/attributes/test_dict.py            | 10 +++--
 tests/data/attributes/test_list.py            | 12 +++--
 tests/data/base/test_attribute.py             |  4 +-
 tests/devices/qubit/test_measure.py           |  8 ++--
 tests/devices/qubit/test_sampling.py          | 14 +++---
 .../test_qutrit_mixed_sampling.py             |  2 +-
 tests/devices/test_default_qubit_legacy.py    | 44 ++++++++++---------
 tests/devices/test_qubit_device.py            |  6 +--
 tests/measurements/test_counts.py             |  6 +--
 .../test_subroutines/test_prepselprep.py      |  9 ++--
 17 files changed, 88 insertions(+), 60 deletions(-)

diff --git a/.github/workflows/install_deps/action.yml b/.github/workflows/install_deps/action.yml
index 809358ce745..99b77dc8157 100644
--- a/.github/workflows/install_deps/action.yml
+++ b/.github/workflows/install_deps/action.yml
@@ -15,7 +15,7 @@ inputs:
   jax_version:
     description: The version of JAX to install for any job that requires JAX
     required: false
-    default: 0.4.23
+    default: '0.4.23'
   install_tensorflow:
     description: Indicate if TensorFlow should be installed or not
     required: false
@@ -23,7 +23,7 @@ inputs:
   tensorflow_version:
     description: The version of TensorFlow to install for any job that requires TensorFlow
     required: false
-    default: 2.16.0
+    default: '2.16.0'
   install_pytorch:
     description: Indicate if PyTorch should be installed or not
     required: false
@@ -31,7 +31,7 @@ inputs:
   pytorch_version:
     description: The version of PyTorch to install for any job that requires PyTorch
     required: false
-    default: 2.3.0
+    default: '2.3.0'
   install_pennylane_lightning_master:
     description: Indicate if PennyLane-Lightning should be installed from the master branch
     required: false
diff --git a/.gitignore b/.gitignore
index fc69a5281fd..d9da3038c69 100644
--- a/.gitignore
+++ b/.gitignore
@@ -27,3 +27,5 @@ config.toml
 qml_debug.log
 datasets/*
 .benchmarks/*
+*.h5
+*.hdf5
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 0a93727855d..506504f631b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -6,6 +6,9 @@
  
 <h3>Improvements 🛠</h3>
 
+* PennyLane is now compatible with NumPy 2.0.
+   [(#6061)](https://github.com/PennyLaneAI/pennylane/pull/6061)
+
 * `qml.qchem.excitations` now optionally returns fermionic operators.
    [(#6171)](https://github.com/PennyLaneAI/pennylane/pull/6171)
 
@@ -124,6 +127,8 @@ This release contains contributions from (in alphabetical order):
 Guillermo Alonso,
 Utkarsh Azad,
 Lillian M. A. Frederiksen,
+Pietropaolo Frisoni,
+Emiliano Godinez,
 Christina Lee,
 William Maxwell,
 Lee J. O'Riordan,
diff --git a/pennylane/numpy/random.py b/pennylane/numpy/random.py
index eae1511cf4f..12a00e798a5 100644
--- a/pennylane/numpy/random.py
+++ b/pennylane/numpy/random.py
@@ -16,9 +16,10 @@
 it works with the PennyLane :class:`~.tensor` class.
 """
 
-from autograd.numpy import random as _random
+# isort: skip_file
 from numpy import __version__ as np_version
 from numpy.random import MT19937, PCG64, SFC64, Philox  # pylint: disable=unused-import
+from autograd.numpy import random as _random
 from packaging.specifiers import SpecifierSet
 from packaging.version import Version
 
@@ -26,8 +27,8 @@
 
 wrap_arrays(_random.__dict__, globals())
 
-
 if Version(np_version) in SpecifierSet(">=0.17.0"):
+
     # pylint: disable=too-few-public-methods
     # pylint: disable=missing-class-docstring
     class Generator(_random.Generator):
diff --git a/requirements-ci.txt b/requirements-ci.txt
index 083beaae25f..d552b95e904 100644
--- a/requirements-ci.txt
+++ b/requirements-ci.txt
@@ -1,5 +1,5 @@
 numpy
-scipy<1.13.0
+scipy<=1.13.0
 cvxpy
 cvxopt
 networkx
diff --git a/setup.py b/setup.py
index e13673fb1fa..41ae9775027 100644
--- a/setup.py
+++ b/setup.py
@@ -21,7 +21,7 @@
     version = f.readlines()[-1].split()[-1].strip("\"'")
 
 requirements = [
-    "numpy<2.0",
+    "numpy<=2.0",
     "scipy",
     "networkx",
     "rustworkx>=0.14.0",
diff --git a/tests/data/attributes/operator/test_operator.py b/tests/data/attributes/operator/test_operator.py
index 83e0c658ab6..8079b720f89 100644
--- a/tests/data/attributes/operator/test_operator.py
+++ b/tests/data/attributes/operator/test_operator.py
@@ -174,6 +174,7 @@ def test_value_init(self, attribute_cls, op_in):
         """Test that a DatasetOperator can be value-initialized
         from an operator, and that the deserialized operator
         is equivalent."""
+
         if not qml.operation.active_new_opmath() and isinstance(op_in, qml.ops.LinearCombination):
             op_in = qml.operation.convert_to_legacy_H(op_in)
 
@@ -183,7 +184,8 @@ def test_value_init(self, attribute_cls, op_in):
         assert dset_op.info["py_type"] == get_type_str(type(op_in))
 
         op_out = dset_op.get_value()
-        assert repr(op_out) == repr(op_in)
+        with np.printoptions(legacy="1.21"):
+            assert repr(op_out) == repr(op_in)
         assert op_in.data == op_out.data
 
     @pytest.mark.parametrize(
@@ -199,6 +201,7 @@ def test_bind_init(self, attribute_cls, op_in):
         """Test that a DatasetOperator can be bind-initialized
         from an operator, and that the deserialized operator
         is equivalent."""
+
         if not qml.operation.active_new_opmath() and isinstance(op_in, qml.ops.LinearCombination):
             op_in = qml.operation.convert_to_legacy_H(op_in)
 
@@ -210,10 +213,12 @@ def test_bind_init(self, attribute_cls, op_in):
         assert dset_op.info["py_type"] == get_type_str(type(op_in))
 
         op_out = dset_op.get_value()
-        assert repr(op_out) == repr(op_in)
+        with np.printoptions(legacy="1.21"):
+            assert repr(op_out) == repr(op_in)
         assert op_in.data == op_out.data
         assert op_in.wires == op_out.wires
-        assert repr(op_in) == repr(op_out)
+        with np.printoptions(legacy="1.21"):
+            assert repr(op_in) == repr(op_out)
 
 
 @pytest.mark.parametrize("attribute_cls", [DatasetOperator, DatasetPyTree])
diff --git a/tests/data/attributes/test_dict.py b/tests/data/attributes/test_dict.py
index 6bf6e202fd6..3e3a2a9d4b2 100644
--- a/tests/data/attributes/test_dict.py
+++ b/tests/data/attributes/test_dict.py
@@ -15,6 +15,7 @@
 Tests for the ``DatasetDict`` attribute type.
 """
 
+import numpy as np
 import pytest
 
 from pennylane.data.attributes import DatasetDict
@@ -45,7 +46,8 @@ def test_value_init(self, value):
         assert dset_dict.info.py_type == "dict"
         assert dset_dict.bind.keys() == value.keys()
         assert len(dset_dict) == len(value)
-        assert repr(value) == repr(dset_dict)
+        with np.printoptions(legacy="1.21"):
+            assert repr(value) == repr(dset_dict)
 
     @pytest.mark.parametrize(
         "value", [{"a": 1, "b": 2}, {}, {"a": 1, "b": {"x": "y", "z": [1, 2]}}]
@@ -93,7 +95,8 @@ def test_copy(self, value):
 
         assert builtin_dict.keys() == value.keys()
         assert len(builtin_dict) == len(value)
-        assert repr(builtin_dict) == repr(value)
+        with np.printoptions(legacy="1.21"):
+            assert repr(builtin_dict) == repr(value)
 
     @pytest.mark.parametrize(
         "value", [{"a": 1, "b": 2}, {}, {"a": 1, "b": {"x": "y", "z": [1, 2]}}]
@@ -121,4 +124,5 @@ def test_equality_same_length(self):
     )
     def test_string_conversion(self, value):
         dset_dict = DatasetDict(value)
-        assert str(dset_dict) == str(value)
+        with np.printoptions(legacy="1.21"):
+            assert str(dset_dict) == str(value)
diff --git a/tests/data/attributes/test_list.py b/tests/data/attributes/test_list.py
index eef27057616..2f4c937d178 100644
--- a/tests/data/attributes/test_list.py
+++ b/tests/data/attributes/test_list.py
@@ -18,6 +18,7 @@
 
 from itertools import combinations
 
+import numpy as np
 import pytest
 
 from pennylane.data import DatasetList
@@ -56,8 +57,9 @@ def test_value_init(self, input_type, value):
 
         lst = DatasetList(input_type(value))
         assert lst == value
-        assert repr(lst) == repr(value)
         assert len(lst) == len(value)
+        with np.printoptions(legacy="1.21"):
+            assert repr(lst) == repr(value)
 
     @pytest.mark.parametrize("input_type", (list, tuple))
     @pytest.mark.parametrize("value", [[], [1], [1, 2, 3], ["a", "b", "c"], [{"a": 1}]])
@@ -148,12 +150,14 @@ def test_setitem_out_of_range(self, index):
     @pytest.mark.parametrize("value", [[], [1], [1, 2, 3], ["a", "b", "c"], [{"a": 1}]])
     def test_copy(self, input_type, value):
         """Test that a `DatasetList` can be copied."""
+
         ds = DatasetList(input_type(value))
         ds_copy = ds.copy()
 
         assert ds_copy == value
-        assert repr(ds_copy) == repr(value)
         assert len(ds_copy) == len(value)
+        with np.printoptions(legacy="1.21"):
+            assert repr(ds_copy) == repr(value)
 
     @pytest.mark.parametrize("input_type", (list, tuple))
     @pytest.mark.parametrize("value", [[], [1], [1, 2, 3], ["a", "b", "c"], [{"a": 1}]])
@@ -169,8 +173,10 @@ def test_equality(self, input_type, value):
     @pytest.mark.parametrize("value", [[], [1], [1, 2, 3], ["a", "b", "c"], [{"a": 1}]])
     def test_string_conversion(self, value):
         """Test that a `DatasetList` is converted to a string correctly."""
+
         dset_dict = DatasetList(value)
-        assert str(dset_dict) == str(value)
+        with np.printoptions(legacy="1.21"):
+            assert str(dset_dict) == str(value)
 
     @pytest.mark.parametrize("value", [[1], [1, 2, 3], ["a", "b", "c"], [{"a": 1}]])
     def test_deleting_elements(self, value):
diff --git a/tests/data/base/test_attribute.py b/tests/data/base/test_attribute.py
index d38249c1672..da500db48e5 100644
--- a/tests/data/base/test_attribute.py
+++ b/tests/data/base/test_attribute.py
@@ -285,8 +285,8 @@ def test_bind_init_from_other_bind(self):
     )
     def test_repr(self, val, attribute_type):
         """Test that __repr__ has the expected format."""
-
-        assert repr(attribute(val)) == f"{attribute_type.__name__}({repr(val)})"
+        with np.printoptions(legacy="1.21"):
+            assert repr(attribute(val)) == f"{attribute_type.__name__}({repr(val)})"
 
     @pytest.mark.parametrize(
         "val",
diff --git a/tests/devices/qubit/test_measure.py b/tests/devices/qubit/test_measure.py
index d0c618311cf..47e4d8c2a31 100644
--- a/tests/devices/qubit/test_measure.py
+++ b/tests/devices/qubit/test_measure.py
@@ -302,7 +302,7 @@ class TestNaNMeasurements:
     def test_nan_float_result(self, mp, interface):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
-        state = qml.math.full((2, 2), np.NaN, like=interface)
+        state = qml.math.full((2, 2), np.nan, like=interface)
         res = measure(mp, state, is_state_batched=False)
 
         assert qml.math.ndim(res) == 0
@@ -339,7 +339,7 @@ def test_nan_float_result(self, mp, interface):
     def test_nan_float_result_jax(self, mp, use_jit):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
-        state = qml.math.full((2, 2), np.NaN, like="jax")
+        state = qml.math.full((2, 2), np.nan, like="jax")
         if use_jit:
             import jax
 
@@ -360,7 +360,7 @@ def test_nan_float_result_jax(self, mp, use_jit):
     def test_nan_probs(self, mp, interface):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
-        state = qml.math.full((2, 2), np.NaN, like=interface)
+        state = qml.math.full((2, 2), np.nan, like=interface)
         res = measure(mp, state, is_state_batched=False)
 
         assert qml.math.shape(res) == (2 ** len(mp.wires),)
@@ -375,7 +375,7 @@ def test_nan_probs(self, mp, interface):
     def test_nan_probs_jax(self, mp, use_jit):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
-        state = qml.math.full((2, 2), np.NaN, like="jax")
+        state = qml.math.full((2, 2), np.nan, like="jax")
         if use_jit:
             import jax
 
diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
index 4174ed63aae..e36c69c26a3 100644
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -591,7 +591,7 @@ def test_only_catch_nan_errors(self, shots):
         mp = qml.expval(qml.PauliZ(0))
         _shots = Shots(shots)
 
-        with pytest.raises(ValueError, match="probabilities do not sum to 1"):
+        with pytest.raises(ValueError, match=r"(?i)probabilities do not sum to 1"):
             _ = measure_with_samples([mp], state, _shots)
 
     @pytest.mark.all_interfaces
@@ -619,7 +619,7 @@ def test_only_catch_nan_errors(self, shots):
     def test_nan_float_result(self, mp, interface, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
-        state = qml.math.full((2, 2), np.NaN, like=interface)
+        state = qml.math.full((2, 2), np.nan, like=interface)
         res = measure_with_samples((mp,), state, _FlexShots(shots), is_state_batched=False)
 
         if not isinstance(shots, list):
@@ -646,7 +646,7 @@ def test_nan_float_result(self, mp, interface, shots):
     def test_nan_samples(self, mp, interface, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
-        state = qml.math.full((2, 2), np.NaN, like=interface)
+        state = qml.math.full((2, 2), np.nan, like=interface)
         res = measure_with_samples((mp,), state, _FlexShots(shots), is_state_batched=False)
 
         if not isinstance(shots, list):
@@ -672,7 +672,7 @@ def test_nan_samples(self, mp, interface, shots):
     def test_nan_classical_shadows(self, interface, shots):
         """Test that classical_shadows returns an empty array when the state has
         NaN values"""
-        state = qml.math.full((2, 2), np.NaN, like=interface)
+        state = qml.math.full((2, 2), np.nan, like=interface)
         res = measure_with_samples(
             (qml.classical_shadow([0]),), state, _FlexShots(shots), is_state_batched=False
         )
@@ -699,7 +699,7 @@ def test_nan_classical_shadows(self, interface, shots):
     def test_nan_shadow_expval(self, H, interface, shots):
         """Test that shadow_expval returns an empty array when the state has
         NaN values"""
-        state = qml.math.full((2, 2), np.NaN, like=interface)
+        state = qml.math.full((2, 2), np.nan, like=interface)
         res = measure_with_samples(
             (qml.shadow_expval(H),), state, _FlexShots(shots), is_state_batched=False
         )
@@ -757,7 +757,7 @@ def test_sample_state_renorm_error(self, interface):
         """Test that renormalization does not occur if the error is too large."""
 
         state = qml.math.array(two_qubit_state_not_normalized, like=interface)
-        with pytest.raises(ValueError, match="probabilities do not sum to 1"):
+        with pytest.raises(ValueError, match=r"(?i)probabilities do not sum to 1"):
             _ = sample_state(state, 10)
 
     @pytest.mark.all_interfaces
@@ -775,7 +775,7 @@ def test_sample_batched_state_renorm_error(self, interface):
         """Test that renormalization does not occur if the error is too large."""
 
         state = qml.math.array(batched_state_not_normalized, like=interface)
-        with pytest.raises(ValueError, match="probabilities do not sum to 1"):
+        with pytest.raises(ValueError, match=r"(?i)probabilities do not sum to 1"):
             _ = sample_state(state, 10, is_state_batched=True)
 
 
diff --git a/tests/devices/qutrit_mixed/test_qutrit_mixed_sampling.py b/tests/devices/qutrit_mixed/test_qutrit_mixed_sampling.py
index eb3383ed5a6..ecd2fbbcca8 100644
--- a/tests/devices/qutrit_mixed/test_qutrit_mixed_sampling.py
+++ b/tests/devices/qutrit_mixed/test_qutrit_mixed_sampling.py
@@ -402,7 +402,7 @@ def test_only_catch_nan_errors(self, shots):
         mp = qml.sample(wires=range(2))
         _shots = Shots(shots)
 
-        with pytest.raises(ValueError, match="probabilities do not sum to 1"):
+        with pytest.raises(ValueError, match=r"(?i)probabilities do not sum to 1"):
             _ = measure_with_samples(mp, state, _shots)
 
     @pytest.mark.parametrize("mp", [qml.probs(0), qml.probs(op=qml.GellMann(0, 1))])
diff --git a/tests/devices/test_default_qubit_legacy.py b/tests/devices/test_default_qubit_legacy.py
index 11ca082441c..9b67d18b5c6 100644
--- a/tests/devices/test_default_qubit_legacy.py
+++ b/tests/devices/test_default_qubit_legacy.py
@@ -18,6 +18,7 @@
 # pylint: disable=protected-access,cell-var-from-loop
 import cmath
 import math
+from importlib.metadata import version
 
 import pytest
 
@@ -628,7 +629,8 @@ def test_apply_global_phase(self, qubit_device_3_wires, tol, wire, input_state):
         expected_output = np.array(input_state) * np.exp(-1j * phase)
 
         assert np.allclose(qubit_device_3_wires._state, np.array(expected_output), atol=tol, rtol=0)
-        assert qubit_device_3_wires._state.dtype == qubit_device_3_wires.C_DTYPE
+        if version("numpy") < "2.0.0":
+            assert qubit_device_3_wires._state.dtype == qubit_device_3_wires.C_DTYPE
 
     def test_apply_errors_qubit_state_vector(self, qubit_device_2_wires):
         """Test that apply fails for incorrect state preparation, and > 2 qubit gates"""
@@ -650,26 +652,26 @@ def test_apply_errors_qubit_state_vector(self, qubit_device_2_wires):
             )
 
     def test_apply_errors_basis_state(self, qubit_device_2_wires):
-
-        with pytest.raises(
-            ValueError, match=r"Basis state must only consist of 0s and 1s; got \[-0\.2, 4\.2\]"
-        ):
-            qubit_device_2_wires.apply([qml.BasisState(np.array([-0.2, 4.2]), wires=[0, 1])])
-
-        with pytest.raises(
-            ValueError, match=r"State must be of length 1; got length 2 \(state=\[0 1\]\)\."
-        ):
-            qubit_device_2_wires.apply([qml.BasisState(np.array([0, 1]), wires=[0])])
-
-        with pytest.raises(
-            qml.DeviceError,
-            match="Operation BasisState cannot be used after other Operations have already been applied "
-            "on a default.qubit.legacy device.",
-        ):
-            qubit_device_2_wires.reset()
-            qubit_device_2_wires.apply(
-                [qml.RZ(0.5, wires=[0]), qml.BasisState(np.array([1, 1]), wires=[0, 1])]
-            )
+        with np.printoptions(legacy="1.21"):
+            with pytest.raises(
+                ValueError, match=r"Basis state must only consist of 0s and 1s; got \[-0\.2, 4\.2\]"
+            ):
+                qubit_device_2_wires.apply([qml.BasisState(np.array([-0.2, 4.2]), wires=[0, 1])])
+
+            with pytest.raises(
+                ValueError, match=r"State must be of length 1; got length 2 \(state=\[0 1\]\)\."
+            ):
+                qubit_device_2_wires.apply([qml.BasisState(np.array([0, 1]), wires=[0])])
+
+            with pytest.raises(
+                qml.DeviceError,
+                match="Operation BasisState cannot be used after other Operations have already been applied "
+                "on a default.qubit.legacy device.",
+            ):
+                qubit_device_2_wires.reset()
+                qubit_device_2_wires.apply(
+                    [qml.RZ(0.5, wires=[0]), qml.BasisState(np.array([1, 1]), wires=[0, 1])]
+                )
 
 
 class TestExpval:
diff --git a/tests/devices/test_qubit_device.py b/tests/devices/test_qubit_device.py
index 9edc522a408..8f50bb65329 100644
--- a/tests/devices/test_qubit_device.py
+++ b/tests/devices/test_qubit_device.py
@@ -1605,9 +1605,9 @@ def test_samples_to_counts_with_nan(self):
 
         # imitate hardware return with NaNs (requires dtype float)
         samples = qml.math.cast_like(samples, np.array([1.2]))
-        samples[0][0] = np.NaN
-        samples[17][1] = np.NaN
-        samples[850][0] = np.NaN
+        samples[0][0] = np.nan
+        samples[17][1] = np.nan
+        samples[850][0] = np.nan
 
         result = device._samples_to_counts(samples, mp=qml.measurements.CountsMP(), num_wires=2)
 
diff --git a/tests/measurements/test_counts.py b/tests/measurements/test_counts.py
index 3f3badb5c0e..08da35015c9 100644
--- a/tests/measurements/test_counts.py
+++ b/tests/measurements/test_counts.py
@@ -135,9 +135,9 @@ def test_counts_with_nan_samples(self):
         rng = np.random.default_rng(123)
         samples = rng.choice([0, 1], size=(shots, 2)).astype(np.float64)
 
-        samples[0][0] = np.NaN
-        samples[17][1] = np.NaN
-        samples[850][0] = np.NaN
+        samples[0][0] = np.nan
+        samples[17][1] = np.nan
+        samples[850][0] = np.nan
 
         result = qml.counts(wires=[0, 1]).process_samples(samples, wire_order=[0, 1])
 
diff --git a/tests/templates/test_subroutines/test_prepselprep.py b/tests/templates/test_subroutines/test_prepselprep.py
index 95e7f771ef7..7e3b7966c28 100644
--- a/tests/templates/test_subroutines/test_prepselprep.py
+++ b/tests/templates/test_subroutines/test_prepselprep.py
@@ -47,13 +47,16 @@ def test_standard_checks(lcu, control):
 
 def test_repr():
     """Test the repr method."""
+
     lcu = qml.dot([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)])
     control = [0]
 
     op = qml.PrepSelPrep(lcu, control)
-    assert (
-        repr(op) == "PrepSelPrep(coeffs=(0.25, 0.75), ops=(Z(2), X(1) @ X(2)), control=Wires([0]))"
-    )
+    with np.printoptions(legacy="1.21"):
+        assert (
+            repr(op)
+            == "PrepSelPrep(coeffs=(0.25, 0.75), ops=(Z(2), X(1) @ X(2)), control=Wires([0]))"
+        )
 
 
 def _get_new_terms(lcu):