code: def test_lambda_evaluation_jax(diff_method):
      """Test gradient of lambda argument replacement matches
      gradients of individual full evaluations using JAX"""
      import jax
      import jax.numpy as jnp
  
      dev = qml.device("default.qubit", wires=2)
  
      @qml.qnode(dev, diff_method=diff_method)
      def circuit(x, y):
          qml.RX(x, wires=0)
          qml.RY(y[..., 0], wires=0)
          qml.RY(y[..., 1], wires=1)
          return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1))
  
      batch_size = 4
  
      # the first partial argument
      x = jnp.asarray(np.random.uniform(size=()))
  
      # the base value of the second partial argument
      y = jnp.asarray(np.random.uniform(size=2))
  
      # the second partial argument as a function of the inputs
      fn = lambda y0: y + y0 * jnp.ones(2)
  
      # values for the second argument
      y0 = jnp.asarray(np.random.uniform(size=batch_size))
  
      batched_partial_circuit = qml.batch_partial(
          circuit, all_operations=True, x=x, preprocess={"y": fn}
      )
  
      # we could also sum over the batch dimension and use the regular
      # gradient instead of the jacobian, but either works
      grad = jnp.diagonal(jax.jacrev(batched_partial_circuit)(y0))
  
      # check the results against individually executed circuits
      indiv_grad = []
      for y0_indiv in y0:
          grad_wrt_second_arg = jax.grad(circuit, argnums=1)(
              x, y + y0_indiv * np.ones(2)
          )
          grad_wrt_y0 = jnp.sum(grad_wrt_second_arg)
          indiv_grad.append(grad_wrt_y0)
  
      assert np.allclose(grad, indiv_grad)
  
imports: import re
  import pytest
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 13 for calling a global function with example usage: test_lambda_evaluation_jax(diff_method) and returns: qml