class_docstr: Test that the sample_state function works as expected
class_docstr_tok: ["Test","that","the","sample_state","function","works","as","expected"]
docstr: Tests that marginal states can be sampled as expected.
docstr_tok: ["Tests","that","marginal","states","can","be","sampled","as","expected","."]
code: def test_marginal_sample_state(self, wire_order):
      """Tests that marginal states can be sampled as expected."""
      state = np.zeros((2, 2, 2))
      state[:, :, 1] = 0.5  # third wire is always 1
      alltrue_axis = wire_order.index(2)
  
      samples = sample_state(state, 20, wires=wire_order)
      assert all(samples[:, alltrue_axis])
  
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 5 for calling the TestSampleState obj's underlying member method code with example usage: obj.test_marginal_sample_state(wire_order) without return types