class_docstr: Tests for the QutritUnitary class.
class_docstr_tok: ["Tests","for","the","QutritUnitary","class","."]
docstr: Tests that QutritUnitary works with jitting.
docstr_tok: ["Tests","that","QutritUnitary","works","with","jitting","."]
code: def test_qutrit_unitary_jax_jit(self, U, num_wires):
      """Tests that QutritUnitary works with jitting."""
      import jax
      from jax import numpy as jnp
  
      U = jnp.array(U)
      f = lambda m: qml.QutritUnitary(
          m, wires=range(num_wires)
      ).matrix()
      out = jax.jit(f)(U)
      assert qml.math.allclose(
          out, qml.QutritUnitary(U, wires=range(num_wires)).matrix()
      )
  
imports: import pytest
  import numpy
  from scipy.stats import unitary_group
  from gate_data import TSWAP
  import pennylane
  from pennylane.wires import Wires
  from pennylane.operation import DecompositionUndefinedError
label_desc: Use image node_id 9 for calling the TestQutritUnitary obj's underlying member method code with example usage: obj.test_qutrit_unitary_jax_jit(U, num_wires) without return types