class_docstr: Test that JAX integrates with is_abstract
docstr: Test that all tensors are abstract when in JIT mode.
  Note that `jax.grad` has slightly different behaviour, and will
  avoid making abstract tensors for non-differentiable arguments.
docstr_tok: ["Test","that","all","tensors","are","abstract","when","in","JIT","mode",".","Note","that","`","jax.grad","`","has","slightly","different","behaviour",",","and","will","avoid","making","abstract","tensors","for","non-differentiable","arguments","."]
code: def test_jit(self):
      """Test that all tensors are abstract when in JIT mode.
      Note that `jax.grad` has slightly different behaviour, and will
      avoid making abstract tensors for non-differentiable arguments."""
      import jax
      import jax.numpy as jnp
  
      @functools.partial(jax.jit, static_argnums=[2])
      def cost(x, w, w_is_abstract=True):
          y = x**2
          z = jnp.ones([2, 2])
  
          assert qml.math.is_abstract(w) == w_is_abstract
  
          assert qml.math.is_abstract(x)
          assert qml.math.is_abstract(y)
          assert qml.math.is_abstract(z)
  
          return jnp.sum(y)
  
      x = jnp.array([0.5, 0.1])
      w = jnp.array(0.1)
  
      assert not qml.math.is_abstract(w)
      assert not qml.math.is_abstract(x)
  
      res = cost(x, w)
      assert res == 0.26
  
      # NOTE: As of JAX 0.3.1, JAX will trace non-differentiated arguments.
      # As a result, it is no longer possible to assume that non-differentiated
      # arguments will not be abstract.
  
      # # since we only specify argnums=0, w will not be abstract
      # grad = jax.grad(cost, argnums=0)(x, w, w_is_abstract=False)
      # assert np.allclose(grad, 2 * x)
  
      # Otherwise, w will be abstract
      grad = jax.grad(cost, argnums=[0, 1])(x, w, w_is_abstract=True)
      assert np.allclose(grad[0], 2 * x)
  
file_path: pennylane/tests/math/test_is_abstract.py
imports: import functools
  import numpy
  import pytest
  import pennylane
label_desc: Use image node_id 2 for calling the TestJAX obj's underlying member method code with example usage: obj.test_jit() and returns: jnp