code: def jacobian(func, argnum=None, method=None, h=None):
      """Returns the Jacobian as a callable function of vector-valued (functions of) QNodes.
      :func:`~.qjit` and Autograd compatible.
  
      .. note::
  
          When used with :func:`~.qjit`, this function currently only supports the
          Catalyst compiler. See :func:`catalyst.jacobian` for more details.
  
          Please see the Catalyst :doc:`quickstart guide <catalyst:dev/quick_start>`,
          as well as the :doc:`sharp bits and debugging tips <catalyst:dev/sharp_bits>`
          page for an overview of the differences between Catalyst and PennyLane.
  
      Args:
          func (function): A vector-valued Python function or QNode that contains
              a combination of quantum and classical nodes. The output of the computation
              must consist of a single NumPy array (if classical) or a tuple of
              expectation values (if a quantum node)
  
          argnum (int or Sequence[int]): Which argument to take the gradient
              with respect to. If a sequence is given, the Jacobian corresponding
              to all marked inputs and all output elements is returned.
  
          method (str): Specifies the gradient method when used with the :func:`~.qjit`
              decorator. Outside of :func:`~.qjit`, this keyword argument
              has no effect and should not be set. In just-in-time (JIT) mode,
              this can be any of ``["auto", "fd"]``, where:
  
              - ``"auto"`` represents deferring the quantum differentiation to the method
                specified by the QNode, while the classical computation is differentiated
                using traditional auto-diff. Catalyst supports ``"parameter-shift"`` and
                ``"adjoint"`` on internal QNodes. QNodes with ``diff_method="finite-diff"``
                are not supported with ``"auto"``.
  
              - ``"fd"`` represents first-order finite-differences for the entire hybrid
                function.
  
          h (float): The step-size value for the finite-difference (``"fd"``) method within
              :func:`~.qjit` decorated functions. This value has no effect in non-compiled
              functions.
  
      Returns:
          function: the function that returns the Jacobian of the input function with respect to the
          arguments in argnum
  
      .. note::
  
          Due to a limitation in Autograd, this function can only differentiate built-in scalar
          or NumPy array arguments.
  
      For ``argnum=None``, the trainable arguments are inferred dynamically from the arguments
      passed to the function. The returned function takes the same arguments as the original
      function and outputs a ``tuple``. The ``i`` th entry of the ``tuple`` has shape
      ``(*output shape, *shape of args[argnum[i]])``.
  
      If a single trainable argument is inferred, or if a single integer
      is provided as ``argnum``, the tuple is unpacked and its only entry is returned instead.
  
      **Example**
  
      Consider the QNode
  
      .. code-block::
  
          dev = qml.device("default.qubit", wires=2)
  
          @qml.qnode(dev)
          def circuit(weights):
              qml.RX(weights[0, 0, 0], wires=0)
              qml.RY(weights[0, 0, 1], wires=1)
              qml.RZ(weights[1, 0, 2], wires=0)
              return qml.probs()
  
          weights = np.array(
              [[[0.2, 0.9, -1.4]], [[0.5, 0.2, 0.1]]], requires_grad=True
          )
  
      It has a single array-valued QNode argument with shape ``(2, 1, 3)`` and outputs
      the probability of each 2-wire basis state, of which there are ``2**num_wires`` = 4.
      Therefore, the Jacobian of this QNode will be a single array with shape ``(2, 2, 1, 3)``:
  
      >>> qml.jacobian(circuit)(weights).shape
      (4, 2, 1, 3)
  
      On the other hand, consider the following QNode for the same circuit
      structure:
  
      .. code-block::
  
          dev = qml.device("default.qubit", wires=2)
  
          @qml.qnode(dev)
          def circuit(x, y, z):
              qml.RX(x, wires=0)
              qml.RY(y, wires=1)
              qml.RZ(z, wires=0)
              return tuple(qml.expval(qml.PauliZ(w)) for w in dev.wires)
  
          x = np.array(0.2, requires_grad=True)
          y = np.array(0.9, requires_grad=True)
          z = np.array(-1.4, requires_grad=True)
  
      It has three scalar QNode arguments and outputs the probability for each of
      the 4 basis states. Consequently, its Jacobian will be a three-tuple of
      arrays with the shape ``(4,)``:
  
      >>> jac = qml.jacobian(circuit)(x, y, z)
      >>> type(jac)
      tuple
      >>> for sub_jac in jac:
      ...     print(sub_jac.shape)
      (4,)
      (4,)
      (4,)
  
      For a more advanced setting of QNode arguments, consider the QNode
  
      .. code-block::
  
          dev = qml.device("default.qubit", wires=3)
  
          @qml.qnode(dev)
          def circuit(x, y):
              qml.RX(x[0], wires=0)
              qml.RY(y[0, 3], wires=1)
              qml.RX(x[1], wires=2)
              return [qml.expval(qml.PauliZ(w)) for w in [0, 1, 2]]
  
          x = np.array([0.1, 0.5], requires_grad=True)
          y = np.array([[-0.3, 1.2, 0.1, 0.9], [-0.2, -3.1, 0.5, -0.7]], requires_grad=True)
  
      If we do not provide ``argnum``, ``qml.jacobian`` will correctly identify both,
      ``x`` and ``y``, as trainable function arguments:
  
      >>> jac = qml.jacobian(circuit)(x, y)
      >>> print(type(jac), len(jac))
      <class 'tuple'> 2
      >>> qml.math.shape(jac[0])
      (8, 2)
      >>> qml.math.shape(jac[1])
      (8, 2, 4)
  
      As we can see, there are two entries in the output, one Jacobian for each
      QNode argument. The shape ``(8, 2)`` of the first Jacobian is the combination
      of the QNode output shape (``(8,)``) and the shape of ``x`` (``(2,)``).
      Similarly, the shape ``(2, 4)`` of ``y`` leads to a Jacobian shape ``(8, 2, 4)``.
  
      Instead we may choose the output to contain only one of the two
      entries by providing an iterable as ``argnum``:
  
      >>> jac = qml.jacobian(circuit, argnum=[1])(x, y)
      >>> print(type(jac), len(jac))
      <class 'tuple'> 1
      >>> qml.math.shape(jac)
      (1, 8, 2, 4)
  
      Here we included the size of the tuple in the shape analysis, corresponding to the
      first dimension of size ``1``.
  
      Finally, we may want to receive the single entry above directly, not as a tuple
      with a single entry. This is done by providing a single integer as ``argnum``
  
      >>> jac = qml.jacobian(circuit, argnum=1)(x, y)
      >>> print(type(jac), len(jac))
      <class 'numpy.ndarray'> 8
      >>> qml.math.shape(jac)
      (8, 2, 4)
  
      As expected, the tuple was unpacked and we directly received the Jacobian of the
      QNode with respect to ``y``.
  
      We can also compute the Jacobian transformation inside a :func:`~.qjit` decorated program:
  
      .. code-block::
  
          dev = qml.device("lightning.qubit", wires=1)
  
          @qml.qjit
          def workflow(x):
              @qml.qnode(dev)
              def circuit(x):
                  qml.RX(np.pi * x[0], wires=0)
                  qml.RY(x[1], wires=0)
                  return qml.probs()
  
              g = qml.jacobian(circuit)
              return g(x)
  
      >>> workflow(np.array([2.0, 1.0]))
      array([[-1.32116540e-07,  1.33781874e-07],
             [-4.20735506e-01,  4.20735506e-01]])
  
      You can further compute the Jacobian transformation using other supported differentiation
      methods by :func:`catalyst.jacobian`.
  
      .. code-block::
  
          @qml.qjit
          def workflow(x):
              @qml.qnode(dev)
              def circuit(x):
                  qml.RX(np.pi * x[0], wires=0)
                  qml.RY(x[1], wires=0)
                  return qml.probs()
  
              g = qml.jacobian(circuit, method="fd", h=0.3)
              return g(x)
  
      >>> qml.qjit(workflow)(np.array([2.0, 1.0]))
      array([[-0.37120096, -0.45467246],
              [0.37120096,  0.45467246]])
      """
      # pylint: disable=no-value-for-parameter
  
      if active_jit := compiler.active_compiler():
          available_eps = compiler.AvailableCompilers.names_entrypoints
          ops_loader = available_eps[active_jit]["ops"].load()
          return ops_loader.jacobian(
              func, method=method, h=h, argnum=argnum
          )
  
      if method or h:
          raise ValueError(
              f"Invalid values for 'method={method}' and 'h={h}' in interpreted mode"
          )
  
      def _get_argnum(args):
          """Inspect the arguments for differentiability and return the
          corresponding indices."""
          argnum = []
  
          for idx, arg in enumerate(args):
              trainable = getattr(
                  arg, "requires_grad", None
              ) or isinstance(arg, ArrayBox)
              if trainable:
                  if arg.dtype.name[:3] == "int":
                      raise ValueError(
                          "Autograd does not support differentiation of ints."
                      )
                  argnum.append(idx)
  
          return argnum
  
      def _jacobian_function(*args, **kwargs):
          """Compute the autograd Jacobian.
  
          This wrapper function is returned to the user instead of autograd.jacobian,
          so that we can take into account cases where the user computes the
          jacobian function once, but then calls it with arguments that change
          in differentiability.
          """
          if argnum is None:
              # Infer which arguments to consider trainable
              _argnum = _get_argnum(args)
              # Infer whether to unpack from the inferred argnum
              unpack = len(_argnum) == 1
          else:
              # For a single integer as argnum, unpack the Jacobian tuple
              unpack = isinstance(argnum, int)
              _argnum = [argnum] if unpack else argnum
  
          if not _argnum:
              warnings.warn(
                  "Attempted to differentiate a function with no trainable parameters. "
                  "If this is unintended, please add trainable parameters via the "
                  "'requires_grad' attribute or 'argnum' keyword."
              )
          jac = tuple(
              _jacobian(func, arg)(*args, **kwargs) for arg in _argnum
          )
  
          return jac[0] if unpack else jac
  
      return _jacobian_function
  
file_path: pennylane/pennylane/_grad.py
imports: import warnings
  from autograd import jacobian
  from autograd.core import make_vjp
  from autograd.numpy.numpy_boxes import ArrayBox
  from autograd.extend import vspace
  from autograd.wrap_util import unary_to_nary
  from pennylane.compiler import compiler
  from pennylane.compiler.compiler import CompileError
label_desc: Use image node_id 1 for calling a global function with example usage: jacobian(func, argnum, method, h) and returns: _jacobian_function, ops_loader, argnum, unknown