class_docstr: Tests the integration of ``_reconstruct_equ`` and ``_reconstruct_gen`` via
  the full ``reconstruct`` function as well as the differentiability of the
  reconstructed function with respect to their single scalar argument.
docstr: Tests the reconstruction and differentiability with JAX.
docstr_tok: ["Tests","the","reconstruction","and","differentiability","with","JAX","."]
code: def test_differentiability_jax(
      self,
      qnode,
      params,
      ids,
      nums_frequency,
      spectra,
      shifts,
      exp_calls,
  ):
      """Tests the reconstruction and differentiability with JAX."""
      import jax
  
      params = tuple(jax.numpy.array(par) for par in params)
      qnode = qml.QNode(qnode, dev_1, interface="jax")
      with qml.Tracker(qnode.device) as tracker:
          recons = reconstruct(
              qnode, ids, nums_frequency, spectra, shifts
          )(*params)
      assert tracker.totals["executions"] == exp_calls
      arg_names = list(signature(qnode.func).parameters.keys())
      for outer_key in recons:
          outer_key_num = arg_names.index(outer_key)
          for inner_key, rec in recons[outer_key].items():
              x0 = params[outer_key_num]
              if not pnp.isscalar(x0):
                  x0 = x0[inner_key]
                  shift_vec = qml.math.zeros_like(params[outer_key_num])
                  shift_vec = qml.math.scatter_element_add(
                      shift_vec, inner_key, 1.0
                  )
              shift_vec = (
                  1.0
                  if pnp.isscalar(params[outer_key_num])
                  else shift_vec
              )
              mask = (
                  0.0
                  if pnp.isscalar(params[outer_key_num])
                  else pnp.ones(qml.math.shape(params[outer_key_num]))
                  - shift_vec
              )
              univariate = lambda x: qnode(
                  *params[:outer_key_num],
                  params[outer_key_num] * mask + x * shift_vec,
                  *params[outer_key_num + 1 :],
              )
              exp_qnode_grad = jax.grad(qnode, argnums=outer_key_num)
              exp_grad = jax.grad(univariate)
              grad = jax.grad(rec)
              assert np.isclose(
                  grad(x0), exp_qnode_grad(*params)[inner_key]
              )
              assert np.isclose(grad(x0 + 0.1), exp_grad(x0 + 0.1))
              assert fun_close(grad, exp_grad, samples=3)
  
imports: from inspect import signature
  from itertools import chain
  from functools import reduce
  import pytest
  import numpy
  import pennylane
  from pennylane import numpy
  from pennylane.fourier.reconstruct import _reconstruct_equ, _reconstruct_gen, _prepare_jobs, reconstruct
  from pennylane.fourier.utils import join_spectra
label_desc: Use image node_id 3 for calling the TestReconstruct obj's underlying member method code with example usage: obj.test_differentiability_jax(qnode, params, ids, nums_frequency, spectra, shifts, exp_calls) without return types