From 52e430b06596fb6c4b7e6c4a3df848fffa6dcd67 Mon Sep 17 00:00:00 2001
From: Astral Cai <astral.cai@xanadu.ai>
Date: Thu, 19 Sep 2024 10:47:05 -0400
Subject: [PATCH] Add check_generator to assert_valid

---
 pennylane/operation.py                  |  2 +-
 pennylane/ops/functions/assert_valid.py | 20 ++++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/pennylane/operation.py b/pennylane/operation.py
index 77f989e6383..f5337e24101 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -1449,7 +1449,7 @@ def generator(self):  # pylint: disable=no-self-use
           0.5 * Y(0) + Z(0) @ X(1)
 
         The generator may also be provided in the form of a dense or sparse Hamiltonian
-        (using :class:`.Hermitian` and :class:`.SparseHamiltonian` respectively).
+        (using :class:`.Hamiltonian` and :class:`.SparseHamiltonian` respectively).
 
         The default value to return is ``None``, indicating that the operation has
         no defined generator.
diff --git a/pennylane/ops/functions/assert_valid.py b/pennylane/ops/functions/assert_valid.py
index e3795ff6608..e28e4936556 100644
--- a/pennylane/ops/functions/assert_valid.py
+++ b/pennylane/ops/functions/assert_valid.py
@@ -169,6 +169,25 @@ def _check_eigendecomposition(op):
         assert qml.math.allclose(decomp_mat, original_mat), failure_comment
 
 
+def _check_generator(op):
+    """Checks that if an operator has a generator, it does."""
+
+    if op.has_generator:
+        gen = op.generator()
+        assert isinstance(gen, qml.Hamiltonian) or isinstance(gen, qml.ops.SparseHamiltonian)
+        new_op = qml.exp(gen, 1j * op.data[0])
+        assert qml.math.allclose(
+            qml.matrix(op, wire_order=op.wires), qml.matrix(new_op, wire_order=op.wires)
+        )
+    else:
+        failure_comment = (
+            "If has_generator is False, the matrix method must raise a ``GeneratorUndefinedError``."
+        )
+        _assert_error_raised(
+            op.matrix, qml.operation.GeneratorUndefinedError, failure_comment=failure_comment
+        )()
+
+
 def _check_copy(op):
     """Check that copies and deep copies give identical objects."""
     copied_op = copy.copy(op)
@@ -333,4 +352,5 @@ def __init__(self, wires):
     _check_matrix(op)
     _check_matrix_matches_decomp(op)
     _check_eigendecomposition(op)
+    _check_generator(op)
     _check_capture(op)