class_docstr: Tests kernel matrix computations.
class_docstr_tok: ["Tests","kernel","matrix","computations","."]
docstr: Test differentiability of the kernel matrix methods with JAX.
docstr_tok: ["Test","differentiability","of","the","kernel","matrix","methods","with","JAX","."]
code: def test_jax(self):
      """Test differentiability of the kernel matrix methods with JAX."""
      import jax
  
      jax.config.update("jax_enable_x64", True)
      jnp = jax.numpy
  
      X1 = jnp.array(self.X1)
      X2 = jnp.array(self.X2)
      K1 = kern.square_kernel_matrix(X1, _diffable_kernel, False)
      K2 = kern.kernel_matrix(X1, X2, _diffable_kernel)
      K3 = kern.square_kernel_matrix(X1, _diffable_kernel, True)
  
      assert qml.math.allclose(K1, self.expected_K1)
      assert qml.math.allclose(K2, self.expected_K2)
      assert qml.math.allclose(K3, self.expected_K3)
  
      dK1 = jax.jacobian(kern.square_kernel_matrix, argnums=0)(
          X1, _diffable_kernel, False
      )
      assert qml.math.allclose(dK1, self.expected_dK1)
      dK2 = jax.jacobian(kern.kernel_matrix, argnums=(0, 1))(
          X1, X2, _diffable_kernel
      )
      assert qml.math.allclose(dK2[0], self.expected_dK2[0])
      assert qml.math.allclose(dK2[1], self.expected_dK2[1])
      dK3 = jax.jacobian(kern.square_kernel_matrix, argnums=0)(
          X1, _diffable_kernel, True
      )
      assert qml.math.allclose(dK3, self.expected_dK3)
  
imports: import math
  import sys
  from functools import partial
  import pytest
  import numpy
  import pennylane
  import pennylane.kernels
  from pennylane import numpy
label_desc: Use image node_id 5 for calling the TestKernelMatrix obj's underlying member method code with example usage: obj.test_jax() without return types