code: def _batch_jacobian(y, x, tape):
      """Computes a Jacobian w.r.t. last dimensions of y and x."""
      # y and x must have the same batch dimensions.
      # For input shapes (b, dy), (b, dx) yields shape (b, dy, dx).
      d = y.shape.as_list()[-1]
      if d is None:
          raise ValueError(
              "Last dimension of state Tensors must be known."
          )
      grads = []
      for i in range(d):
          w = tf.broadcast_to(tf.one_hot(i, d, dtype=y.dtype), y.shape)
          # We must use tf.UnconnectedGradients.ZERO here and below, because some
          # state components may legitimately not depend on each other or some of the
          # params.
          grad = tape.gradient(
              y,
              x,
              output_gradients=w,
              unconnected_gradients=tf.UnconnectedGradients.ZERO,
          )
          grads.append(grad)
      return tf.stack(grads, axis=-2)
  
file_path: tf-quant-finance/tf_quant_finance/math/custom_loops.py
imports: import tensorflow.compat.v2
  from tf_quant_finance import utils
label_desc: Use image node_id 4 for calling a global function with example usage: _batch_jacobian(y, x, tape) and returns: tf