code: def _make_unit_jacobian(initial_state, params):
      """Creates a unit Jacobian matrix."""
      n = len(initial_state)
      d = [initial_state[i].shape.as_list()[-1] for i in range(n)]
      if None in d:
          raise ValueError(
              "Last dimensions of initial_state Tensors must be known."
          )
      p = len(params)
      dtype = initial_state[0].dtype
  
      def make_js_block(i, j):
          shape = initial_state[i].shape.concatenate((d[j],))
          if i != j:
              return tf.zeros(shape, dtype=dtype)
          eye = tf.eye(d[i], dtype=dtype)
          return tf.broadcast_to(eye, shape)
  
      def make_jp_block(i, j):
          del j
          shape = initial_state[i].shape.concatenate((1,))
          return tf.zeros(shape, dtype=dtype)
  
      js = [[make_js_block(i, j) for j in range(n)] for i in range(n)]
      jp = [[make_jp_block(i, j) for j in range(p)] for i in range(n)]
      return js, jp
  
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 2 for calling a global function with example usage: _make_unit_jacobian(initial_state, params) and returns: js,  jp, tf, tf, tf