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 sample_state calls _sample_state_jax if the seed is a JAX PRNG key
docstr_tok: ["Tests","that","sample_state","calls","_sample_state_jax","if","the","seed","is","a","JAX","PRNG","key"]
code: def test_prng_key_as_seed_uses_sample_state_jax(self, mocker):
      """Tests that sample_state calls _sample_state_jax if the seed is a JAX PRNG key"""
      import jax
  
      jax.config.update("jax_enable_x64", True)
  
      spy = mocker.spy(qml.devices.qubit.sampling, "_sample_state_jax")
      state = qml.math.array(two_qubit_state, like="jax")
  
      # prng_key specified, should call _sample_state_jax
      _ = sample_state(state, 10, prng_key=jax.random.PRNGKey(15))
      # prng_key defaults to None, should NOT call _sample_state_jax
      _ = sample_state(state, 10, rng=15)
  
      spy.assert_called_once()
  
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 2 for calling the TestSampleState obj's underlying member method code with example usage: obj.test_prng_key_as_seed_uses_sample_state_jax(mocker) without return types