class_docstr: Test that measurements work and use _sample_state_jax when the state has a batch dim
  and a PRNG key is provided
class_docstr_tok: ["Test","that","measurements","work","and","use","_sample_state_jax","when","the","state","has","a","batch","dim","and","a","PRNG","key","is","provided"]
docstr: Test that broadcasting works for qml.sample and single shots
docstr_tok: ["Test","that","broadcasting","works","for","qml.sample","and","single","shots"]
code: def test_sample_measure(self, mocker):
      """Test that broadcasting works for qml.sample and single shots"""
      import jax
  
      jax.config.update("jax_enable_x64", True)
  
      spy = mocker.spy(qml.devices.qubit.sampling, "_sample_state_jax")
  
      rng = np.random.default_rng(123)
      shots = qml.measurements.Shots(100)
  
      state = [
          np.array([[0, 0], [0, 1]]),
          np.array([[1, 0], [1, 0]]) / np.sqrt(2),
          np.array([[1, 1], [1, 1]]) / 2,
      ]
      state = np.stack(state)
  
      measurement = qml.sample(wires=[0, 1])
      res = measure_with_samples(
          [measurement],
          state,
          shots,
          is_state_batched=True,
          rng=rng,
          prng_key=jax.random.PRNGKey(184),
      )[0]
  
      spy.assert_called()
  
      assert res.shape == (3, shots.total_shots, 2)
      assert res.dtype == np.int64
  
      # convert to numpy array because prng_key -> JAX -> ArrayImpl -> angry vanilla numpy below
      res = [np.array(r) for r in res]
  
      # first batch of samples is always |11>
      assert np.all(res[0] == 1)
  
      # second batch of samples is either |00> or |10>
      assert np.all(np.logical_or(res[1] == [0, 0], res[1] == [1, 0]))
  
      # third batch of samples can be any of |00>, |01>, |10>, or |11>
      assert np.all(np.logical_or(res[2] == 0, res[2] == 1))
  
imports: from random import shuffle
  import pytest
  import numpy
  import pennylane
  from pennylane.devices.qubit import simulate
  from pennylane.devices.qubit.simulate import _FlexShots
  from pennylane.devices.qubit import sample_state, measure_with_samples
  from pennylane.devices.qubit.sampling import _sample_state_jax
  from pennylane.measurements import Shots
label_desc: Use image node_id 1 for calling the TestBroadcastingPRNG obj's underlying member method code with example usage: obj.test_sample_measure(mocker) without return types