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 the returned samples are as expected when explicitly calling _sample_state_jax.
docstr_tok: ["Tests","that","the","returned","samples","are","as","expected","when","explicitly","calling","_sample_state_jax","."]
code: def test_sample_state_jax(self):
      """Tests that the returned samples are as expected when explicitly calling _sample_state_jax."""
      import jax
  
      state = qml.math.array(two_qubit_state, like="jax")
  
      samples = _sample_state_jax(
          state, 10, prng_key=jax.random.PRNGKey(84)
      )
  
      assert samples.shape == (10, 2)
      assert samples.dtype == np.int64
      assert all(
          qml.math.allequal(s, [0, 1]) or qml.math.allequal(s, [1, 0])
          for s in samples
      )
  
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 TestSampleState obj's underlying member method code with example usage: obj.test_sample_state_jax() without return types