code: def jacobian(
      func,
      x,
      unconnected_gradients=None,
      parallel_iterations=None,
      experimental_use_pfor=True,
      name=None,
  ):
      """Computes the jacobian of `func` wrt to `x`.
  
      Args:
        func: Python callable accepting one `Tensor` of shape of `x` and returning
          a `Tensor` of any shape. The function whose jacobian is to be computed.
        x: A `Tensor` with respect to which the gradient is to be computed.
        unconnected_gradients: An enum `tf.UnconnectedGradients` which specifies
          the gradient value returned when the given input tensors are
          unconnected. Default value: `None`, which maps to
          `tf.UnconnectedGradients.NONE`.
        parallel_iterations: A knob to control how many iterations are dispatched
          in parallel. This knob can be used to control the total memory usage.
        experimental_use_pfor: If true, uses pfor for computing the Jacobian.
          Else uses a tf.while_loop.
        name: Python `str` name prefixed to ops created by this function.
          Default value: `None` (i.e., 'jacobian').
  
      Returns:
        A `Tensor` with the gradient of `y` wrt each of `x`.
      """
      unconnected_gradients = (
          unconnected_gradients or tf.UnconnectedGradients.NONE
      )
      x, is_x_batch_size = _prepare_args(x)
      with tf.name_scope(name or "jacobian"):
          if not callable(func):
              raise ValueError(
                  "`func` should be a callable in eager mode or "
                  "when `tf.GradientTape` is used."
              )
          with tf.GradientTape() as tape:
              tape.watch(x)
              y = func(x)
          jac = tape.batch_jacobian(
              y,
              x,
              unconnected_gradients=unconnected_gradients,
              parallel_iterations=parallel_iterations,
              experimental_use_pfor=experimental_use_pfor,
          )
  
          if is_x_batch_size:
              return jac
  
          return jac[0]
  
imports: import tensorflow.compat.v2
label_desc: Use image node_id 1 for calling a global function with example usage: jacobian(func, x, unconnected_gradients, parallel_iterations, experimental_use_pfor, name) and returns: jac, jac