class_docstr: Tests for the probs_to_unitary function
docstr: Test if the correct unitary is returned for fixed input examples using JAX-JIT.
  A correct unitary has its first column equal to the square root of the distribution
  and satisfies U @ U.T = U.T @ U = I.
docstr_tok: ["Test","if","the","correct","unitary","is","returned","for","fixed","input","examples","using","JAX-JIT",".","A","correct","unitary","has","its","first","column","equal","to","the","square","root","of","the","distribution","and","satisfies","U","@","U.T","=","U.T","@","U","=","I","."]
code: def test_fixed_examples_jax_jit(self, p):
      """Test if the correct unitary is returned for fixed input examples using JAX-JIT.
      A correct unitary has its first column equal to the square root of the distribution
      and satisfies U @ U.T = U.T @ U = I."""
      import jax
      from jax import numpy as jnp
  
      unitary = jax.jit(probs_to_unitary)(jnp.array(p))
      assert jnp.allclose(np.sqrt(p), unitary[:, 0])
      assert jnp.allclose(
          unitary @ unitary.T, np.eye(len(unitary)), atol=1e-7
      )
      assert jnp.allclose(
          unitary.T @ unitary, np.eye(len(unitary)), atol=1e-7
      )
  
imports: import numpy
  import pytest
  from scipy.stats import norm
  import pennylane
  from pennylane.templates.subroutines.qmc import QuantumMonteCarlo, _make_V, _make_Z, func_to_unitary, make_Q, probs_to_unitary
  from pennylane.wires import Wires
label_desc: Use image node_id 4 for calling the TestProbsToUnitary obj's underlying member method code with example usage: obj.test_fixed_examples_jax_jit(p) without return types