class_docstr: Tests for the probs function
class_docstr_tok: ["Tests","for","the","probs","function"]
docstr: Test the rotated probability with different wires and rotating operations.
docstr_tok: ["Test","the","rotated","probability","with","different","wires","and","rotating","operations","."]
code: def test_operation_prob(
      self, operation, wire, init_state, tol, mocker
  ):
      "Test the rotated probability with different wires and rotating operations."
      # pylint:disable=too-many-arguments
      dev = qml.device("default.qubit.legacy", wires=4)
      spy = mocker.spy(qml.QubitDevice, "probability")
      state = init_state(4)
  
      @qml.qnode(dev)
      def circuit():
          qml.StatePrep(state, wires=list(range(4)))
          qml.PauliX(wires=0)
          qml.PauliZ(wires=1)
          qml.PauliY(wires=2)
          qml.PauliZ(wires=3)
          return qml.probs(op=operation(wires=wire))
  
      res = circuit()
  
      def circuit_rotated():
          qml.PauliX(wires=0)
          qml.PauliZ(wires=1)
          qml.PauliY(wires=2)
          qml.PauliZ(wires=3)
          operation(wires=wire).diagonalizing_gates()
  
      matrix = qml.matrix(circuit_rotated)()
      state = np.dot(matrix, state)
      expected = np.reshape(np.abs(state) ** 2, [2] * 4)
  
      if wire == 0:
          expected = np.einsum("ijkl->i", expected).flatten()
      elif wire == 1:
          expected = np.einsum("ijkl->j", expected).flatten()
      elif wire == 2:
          expected = np.einsum("ijkl->k", expected).flatten()
      elif wire == 3:
          expected = np.einsum("ijkl->l", expected).flatten()
  
      assert np.allclose(res, expected, atol=tol, rtol=0)
  
      custom_measurement_process(dev, spy)
  
imports: import numpy
  import pytest
  import pennylane
  from pennylane import numpy
  from pennylane.measurements import ProbabilityMP, Shots
  from pennylane.devices.qubit.measure import flatten_state
label_desc: Use image node_id 18 for calling the TestProbs obj's underlying member method code with example usage: obj.test_operation_prob(operation, wire, init_state, tol, mocker) and returns: qml