class_docstr: Test that measurements work when the state has a batch dim
class_docstr_tok: ["Test","that","measurements","work","when","the","state","has","a","batch","dim"]
docstr: Test that broadcasting works for qml.sample and shot vectors
docstr_tok: ["Test","that","broadcasting","works","for","qml.sample","and","shot","vectors"]
code: def test_sample_measure_shot_vector(self, shots):
      """Test that broadcasting works for qml.sample and shot vectors"""
  
      rng = np.random.default_rng(123)
      shots = qml.measurements.Shots(shots)
  
      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
      )
  
      assert isinstance(res, tuple)
      assert len(res) == shots.num_copies
  
      for s, r in zip(shots, res):
          assert isinstance(r, tuple)
          assert len(r) == 1
          r = r[0]
  
          assert r.shape == (3, s, 2)
          assert r.dtype == np.int64
  
          # first batch of samples is always |11>
          assert np.all(r[0] == 1)
  
          # second batch of samples is either |00> or |10>
          assert np.all(np.logical_or(r[1] == [0, 0], r[1] == [1, 0]))
  
          # third batch of samples can be any of |00>, |01>, |10>, or |11>
          assert np.all(np.logical_or(r[2] == 0, r[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 3 for calling the TestBroadcasting obj's underlying member method code with example usage: obj.test_sample_measure_shot_vector(shots) without return types