class_docstr: Test that the sample_state function works as expected
class_docstr_tok: ["Test","that","the","sample_state","function","works","as","expected"]
docstr: Test that setting the seed as a JAX PRNG key determines the results for _sample_state_jax
docstr_tok: ["Test","that","setting","the","seed","as","a","JAX","PRNG","key","determines","the","results","for","_sample_state_jax"]
code: def test_prng_key_determines_sample_state_jax_results(self):
      """Test that setting the seed as a JAX PRNG key determines the results for _sample_state_jax"""
      import jax
  
      state = qml.math.array(two_qubit_state, like="jax")
  
      samples = _sample_state_jax(
          state, shots=10, prng_key=jax.random.PRNGKey(12)
      )
      samples2 = _sample_state_jax(
          state, shots=10, prng_key=jax.random.PRNGKey(12)
      )
      samples3 = _sample_state_jax(
          state, shots=10, prng_key=jax.random.PRNGKey(13)
      )
  
      assert np.all(samples == samples2)
      assert not np.allclose(samples, samples3)
  
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 4 for calling the TestSampleState obj's underlying member method code with example usage: obj.test_prng_key_determines_sample_state_jax_results() without return types