class_docstr: Simulator plugin based on ``"default.qubit.legacy"``, written using jax.
  
  **Short name:** ``default.qubit.jax``
  
  This device provides a pure-state qubit simulator written using jax. As a result, it
  supports classical backpropagation as a means to compute the gradient. This can be faster than
  the parameter-shift rule for analytic quantum gradients when the number of parameters to be
  optimized is large.
  
  To use this device, you will need to install jax:
  
  .. code-block:: console
  
      pip install jax jaxlib
  
  **Example**
  
  The ``default.qubit.jax`` device is designed to be used with end-to-end classical backpropagation
  (``diff_method="backprop"``) with the JAX interface. This is the default method of
  differentiation when creating a QNode with this device.
  
  Using this method, the created QNode is a 'white-box', and is
  tightly integrated with your JAX computation:
  
  >>> dev = qml.device("default.qubit.jax", wires=1)
  >>> @qml.qnode(dev, interface="jax", diff_method="backprop")
  ... def circuit(x):
  ...     qml.RX(x[1], wires=0)
  ...     qml.Rot(x[0], x[1], x[2], wires=0)
  ...     return qml.expval(qml.PauliZ(0))
  >>> weights = jnp.array([0.2, 0.5, 0.1])
  >>> grad_fn = jax.grad(circuit)
  >>> print(grad_fn(weights))
  array([-2.2526717e-01 -1.0086454e+00  1.3877788e-17])
  
  There are a couple of things to keep in mind when using the ``"backprop"``
  differentiation method for QNodes:
  
  * You must use the ``"jax"`` interface for classical backpropagation, as JAX is
    used as the device backend.
  
  .. details::
      :title: Usage Details
  
      JAX does randomness in a special way when compared to NumPy, in that all randomness needs to
      be seeded. While we handle this for you automatically in op-by-op mode, when using ``jax.jit``,
      the automatically generated seed gets constantant compiled.
  
      Example:
  
      .. code-block:: python
  
          dev = qml.device("default.qubit.jax", wires=1, shots=10)
  
          @jax.jit
          @qml.qnode(dev, interface="jax", diff_method="backprop")
          def circuit():
              qml.Hadamard(0)
              return qml.sample(qml.PauliZ(wires=0))
  
          a = circuit()
          b = circuit() # Bad! b will be the exact same samples as a.
  
  
      To fix this, you should wrap your qnode in another function that takes a PRNGKey, and pass
      that in during your device construction.
  
      .. code-block:: python
  
          @jax.jit
          def keyed_circuit(key):
              dev = qml.device("default.qubit.jax", prng_key=key, wires=1, shots=10)
              @qml.qnode(dev, interface="jax", diff_method="backprop")
              def circuit():
                  qml.Hadamard(0)
                  return qml.sample(qml.PauliZ(wires=0))
              return circuit()
  
          key1 = jax.random.PRNGKey(0)
          key2 = jax.random.PRNGKey(1)
          a = keyed_circuit(key1)
          b = keyed_circuit(key2) # b will be different samples now.
  
      Check out the `JAX random documentation <https://jax.readthedocs.io/en/latest/jax.random.html>`__
      for more information.
  
  Args:
      wires (int): The number of wires to initialize the device with.
      shots (None, int): How many times the circuit should be evaluated (or sampled) to estimate
          the expectation values. Defaults to ``None`` if not specified, which means that the device
          returns analytical results.
      analytic (bool): Indicates if the device should calculate expectations
          and variances analytically. In non-analytic mode, the ``diff_method="backprop"``
          QNode differentiation method is not supported and it is recommended to consider
          switching device to ``default.qubit`` and using ``diff_method="parameter-shift"``.
          Or keeping ``default.qubit.jax`` but switching to
          ``diff_method=qml.gradients.stoch_pulse_grad`` for pulse programming.
      prng_key (Optional[jax.random.PRNGKey]): An optional ``jax.random.PRNGKey``. This is the key to the
          pseudo random number generator. If None, a random key will be generated.
docstr: Sample from the computational basis states based on the state
  probability.
  
  This is an auxiliary method to the generate_samples method.
  
  Args:
      number_of_states (int): the number of basis states to sample from
  
  Returns:
      List[int]: the sampled basis states
docstr_tok: ["Sample","from","the","computational","basis","states","based","on","the","state","probability",".","This","is","an","auxiliary","method","to","the","generate_samples","method",".","Args",":","number_of_states","(","int",")",":","the","number","of","basis","states","to","sample","from","Returns",":","List","[","int","]",":","the","sampled","basis","states"]
code: def sample_basis_states(self, number_of_states, state_probability):
      """Sample from the computational basis states based on the state
      probability.
  
      This is an auxiliary method to the generate_samples method.
  
      Args:
          number_of_states (int): the number of basis states to sample from
  
      Returns:
          List[int]: the sampled basis states
      """
      if self.shots is None:
          raise qml.QuantumFunctionError(
              "The number of shots has to be explicitly set on the device "
              "when using sample-based measurements."
          )
  
      shots = self.shots
  
      if self._prng_key is None:
          # Assuming op-by-op, so we'll just make one.
          key = jax.random.PRNGKey(np.random.randint(0, 2**31))
      else:
          key = self._prng_key
      if jnp.ndim(state_probability) == 2:
          # Produce separate keys for each of the probabilities along the broadcasted axis
          keys = []
          for _ in state_probability:
              key, subkey = jax.random.split(key)
              keys.append(subkey)
          return jnp.array(
              [
                  jax.random.choice(
                      _key, number_of_states, shape=(shots,), p=prob
                  )
                  for _key, prob in zip(keys, state_probability)
              ]
          )
      return jax.random.choice(
          key, number_of_states, shape=(shots,), p=state_probability
      )
  
file_path: pennylane/pennylane/devices/default_qubit_jax.py
imports: import numpy
  import pennylane
  from pennylane.devices import DefaultQubitLegacy
  from pennylane.pulse import ParametrizedEvolution
  from pennylane.typing import TensorLike
label_desc: Use image node_id 6 for calling the DefaultQubitJax obj's underlying member method code with example usage: obj.sample_basis_states(number_of_states, state_probability) and returns: jax, jnp