code: def test_riemannian_gradient_omegas_restricted(circuit, hamiltonian):
      """Test that we calculate the (restricted) Riemannian gradient coefficients correctly."""
      # pylint: disable=no-member
      nqubits = max(max(ps.wires) for ps in hamiltonian.ops) + 1
      wires = range(nqubits)
      dev = qml.device("default.qubit", wires=nqubits)
  
      @qml.qnode(dev)
      def get_state():
          circuit()
          return qml.state()
  
      @qml.qnode(dev)
      def test_circuit():
          circuit()
          return qml.expval(hamiltonian)
  
      phi = get_state()
      rho = np.outer(phi, phi.conj())
      hamiltonian_np = hamiltonian.sparse_matrix(
          wire_order=wires
      ).toarray()
      riemannian_gradient_np = (
          hamiltonian_np @ rho - rho @ hamiltonian_np
      )
  
      restriction = qml.Hamiltonian(
          coeffs=[1.0] * 3,
          observables=[
              qml.PauliX(0),
              qml.PauliY(1),
              qml.PauliY(0) @ qml.PauliY(1),
          ],
      )
  
      opt = RiemannianGradientOptimizer(
          circuit=test_circuit, restriction=restriction
      )
      ops = opt.get_su_n_operators(restriction)[0]
      omegas_np = []
      for op in ops:
          op = qml.math.expand_matrix(op.matrix(), op.wires, wires)
          omegas_np.append(1j * np.trace(riemannian_gradient_np @ op))
      omegas = opt.get_omegas()
  
      assert np.allclose(omegas, omegas_np)
  
file_path: pennylane/tests/optimize/test_riemannian_gradient_optimizer.py
imports: import pytest
  from scipy.sparse.linalg import expm
  import numpy
  import pennylane
  from pennylane.optimize import RiemannianGradientOptimizer
label_desc: Use image node_id 5 for calling a global function with example usage: test_riemannian_gradient_omegas_restricted(circuit, hamiltonian) and returns: qml, qml