Update which errors to catch

PennyLaneAI · Sep 19, 2024 · 4a93aa0 · 4a93aa0
1 parent d920143
commit 4a93aa0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pennylane/devices/_qubit_device.py b/pennylane/devices/_qubit_device.py
@@ -1055,7 +1055,7 @@ def vn_entropy(self, wires, log_base):
  """
  try:
  state = self.density_matrix(wires=self.wires)
- except qml.QuantumFunctionError as e: # pragma: no cover
+ except (qml.QuantumFunctionError, NotImplementedError) as e: # pragma: no cover
  raise NotImplementedError(
  f"Cannot compute the Von Neumman entropy with device {self.name} that is not capable of returning the "
  f"state. "
@@ -1080,7 +1080,7 @@ def vn_entanglement_entropy(self, wires0, wires1, log_base):
  """
  try:
  state = self.density_matrix(wires=self.wires)
- except qml.QuantumFunctionError as e: # pragma: no cover
+ except (qml.QuantumFunctionError, NotImplementedError) as e: # pragma: no cover
  raise NotImplementedError(
  f"Cannot compute the Von Neumman entropy with device {self.name} that is not capable of returning the "
  f"state. "
@@ -1112,7 +1112,7 @@ def mutual_info(self, wires0, wires1, log_base):
  """
  try:
  state = self.density_matrix(wires=self.wires)
- except qml.QuantumFunctionError as e: # pragma: no cover
+ except (qml.QuantumFunctionError, NotImplementedError) as e: # pragma: no cover
  raise NotImplementedError(
  f"Cannot compute the mutual information with device {self.name} that is not capable of returning the "
  f"state. "