class_docstr: Tests for the probs function
class_docstr_tok: ["Tests","for","the","probs","function"]
docstr: Test the probability is correct for a batched input.
docstr_tok: ["Test","the","probability","is","correct","for","a","batched","input","."]
code: def test_batch_size(self, mocker, shots):
      """Test the probability is correct for a batched input."""
      dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
      spy = mocker.spy(qml.QubitDevice, "probability")
  
      @qml.qnode(dev)
      def circuit(x):
          qml.RX(x, 0)
          return qml.probs(
              wires=dev.wires
          )  # TODO: Use ``qml.probs()`` when supported
  
      x = np.array([0, np.pi / 2])
      res = circuit(x)
      expected = [[1.0, 0.0], [0.5, 0.5]]
      assert np.allclose(res, expected, atol=0.1, rtol=0.1)
  
      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 12 for calling the TestProbs obj's underlying member method code with example usage: obj.test_batch_size(mocker, shots) and returns: qml