code: def for_loop(
      body_fn, initial_state, params, num_iterations, name=None
  ):
      """A for loop with a custom batched gradient.
  
      A for loop with a custom gradient that in certain cases outperforms the
      tf.while_loop gradient implementation.
  
      This is not a general replacement for tf.while_loop as imposes a number of
      restrictions on the inputs:
      - All tensors in loop state must have the same shape except for the last
      dimension.
      - All dimensions except for the last one are treated as batch dimensions, i.e.
      it is assumed that batch elements don't interact with each other inside the
      loop body.
      - The last dimensions and the number of parameters must be statically known
      and be reasonably small (so that the full Jacobian matrix with respect to them
      can be calculated efficiently, see below).
      - It requires an explicit list of parameters used in the loop body, with
      respect to which one wishes to calculate derivatives. This is different from
      tf.while_loop which automatically deals with tensors captured in the closure.
      - Parameters must be a sequence of zero-dimensional tensors.
      - Arbitrary nested structure of state is not supported, the state must be a
      flat sequence of tensors.
  
      The issue this implementation addresses is the additional while loops created
      by the gradient of `tf.while_loop`. To compute the backward gradient
      (more precisely, the vector-Jacobian product) of a while loop, one needs to
      run the loop "backwards". This implementation avoids creating a second loop by
      calculating the full (batched) Jacobian matrix in the forward pass. It is
      efficient when the non-batched part of the shape of the Jacobian is small.
      This part has size `nd * (nd + p)` where `nd` is the sum of last dimensions of
      tensors in the state and `p` is the number of parameters.
  
      This implementation is suitable for e.g. Monte-Carlo sampling, where the state
      represents a batch of independent paths.
  
      #### Example:
  
      ```python
      x = tf.constant([[3.0, 4.0], [30.0, 40.0]])
      y = tf.constant([[7.0, 8.0], [70.0, 80.0]])
      alpha = tf.constant(2.0)
      beta = tf.constant(1.0)
  
      with tf.GradientTape(persistent=True) as tape:
        tape.watch([alpha, beta])
        def body(i, state):
          x, y = state
          return [x * alpha - beta, y * beta + x]
        x_out, y_out = for_loop(body, [x, y], [alpha, beta], 3)
  
      grad = tape.gradient(y_out, beta)  # Returns tf.Tensor(783.0)
      ```
  
      Args:
        body_fn: A Callable. Accepts an iteration index as a 0-dimensional int32
          tensor and state - a tuple of Tensors of same shape as `initial_state`.
          Should return the output state with the same structure as the input state.
        initial_state: A sequence of Tensors with common batch shape. All dimensions
          except the last are treated as batch shape (i.e. not mixed in loop body).
        params: A list of zero-dimensional Tensors - tensors that `body_fn` uses,
          and with respect to which the differentiation is going to happen.
        num_iterations: A rank 0 or rank 1 integer tensor. If the rank is 1, the
          entries are expected to be unique and ordered and  the output will contain
          results obtained at each iteration number specified in `num_iterations`,
          stacked along the first dimension. E.g. if `initial_state` has shapes
          `(10, 20, 2)` and `(10, 20, 3)`, and `num_iterations = [2, 5, 7, 10]` the
          output is a list of tensors with shapes `(4, 10, 20, 2)` and
          `(4, 10, 20, 3)`.
  
        name: Python str. The name to give to the ops created by this function,
          'for_loop' by default.
  
      Returns:
       A list of Tensors of the same shape as `initial_state`, if `num_iterations`
       is a single integer, or with extra first dimension of size
       `len(num_iterations)` otherwise.
       The outputs are differentiable with respect to `initial_state` and `params`,
       but not any other tensors that are captured by `body_fn`. Differentiating
       with respect to an element of `initial_state` yields a tensor with the same
       shape as that element. Differentiating with respect to one of `params` yields
       a tensor of zero shape. If the output state doesn't depend on the given
       parameter, the tensor will be filled with zeros.
      """
  
      # Implementation explanation.
      #
      # Notation:
      # n - number of state tensors,
      # d - last dim of state (it can be different among the n state tensors, but
      # for illustration we'll assume it's the same),
      # p - number of parameters.
      #
      # The common batch dimensions are omitted below.
      #
      # The Jacobian has the form
      # | Js Jp |
      # | 0  I  |,
      # where
      # Js = d state / d state, shape = (nd, nd),
      # Js = d state / d params, shape = (nd, p),
      # 0 - zero matrix (p, nd),
      # I - unit matrix (p, p).
      #
      # Multiplying two Jacobians yields
      # | Js1 Js2    Js1 Jp2 + Jp1 |
      # |   0              I       |
      #
      # The custom gradient function receives output weights ws with shape = (nd,).
      # We turn them into row vectors ws' with shape(1, nd), multiply by Jacobians,
      # (ws' Js, ws' Jp), yielding shapes (1, nd) and (1, p), and finally squeeze
      # the dimensions to get the desired shapes (nd,), (p,).
      #
      # Js and Jp have block structure themselves:
      #
      # Js =  | Js_11 ... Js_1n|
      #       | ...   ... ...  |
      #       | Js_n1 ... Js_nn|
      #
      # Js =  | Jp_11 ... Jp_1p|
      #       | ...   ... ...  |
      #       | Jp_n1 ... Jp_np|
      #
      # where Js_ij have shape (d, d), Jp_ij have shape (d, 1).
      #
      # Js_ij and Jp_ij are Tensors, and the rest are nested lists. We multiply and
      # add the tensors with TF, and the nested lists - manually with Python loops.
      #
      # Note that we can't concatenate the parameters into a single Tensor to avoid
      # some Python loops, even though it's cheap. There will be no path from
      # tf.concat node to the body_fn output, because the user uses the original
      # (not concatenated) parameters in body_fn.
  
      num_iterations = tf.convert_to_tensor(
          num_iterations, dtype=tf.int32, name="num_iterations"
      )
      num_iterations_shape = num_iterations.shape.as_list()
      if num_iterations_shape is None:
          raise ValueError(
              "Rank of num_iterations must be statically known."
          )
      if len(num_iterations_shape) > 1:
          raise ValueError("Rank of num_iterations must be 0 or 1")
      if len(num_iterations_shape) == 1:
          return _accumulating_for_loop(
              body_fn, initial_state, params, num_iterations, name
          )
  
      with tf.name_scope(name or "for_loop"):
          initial_jac = _make_unit_jacobian(initial_state, params)
          n = len(initial_state)
  
          @tf.custom_gradient
          def inner(*args):
              initial_state, params = args[:n], args[n:]
  
              def while_cond(i, state, jac):
                  del state, jac
                  return i < num_iterations
  
              def while_body(i, state, jac):
                  with tf.GradientTape(persistent=True) as tape:
                      tape.watch(state)
                      tape.watch(params)
                      next_state = tuple(body_fn(i, state))
                  step_jac = _compute_step_jacobian(
                      state, next_state, params, tape
                  )
                  next_jac = _multiply_jacobians(step_jac, jac)
                  return i + 1, next_state, next_jac
  
              loop_vars = (0, initial_state, initial_jac)
  
              _, state, jac = tf.compat.v2.while_loop(
                  while_cond,
                  while_body,
                  loop_vars=loop_vars,
                  maximum_iterations=num_iterations,
              )
  
              def gradient(*ws):
                  # tf.custom_gradient converts any structure of function outputs into a
                  # flat tuple when calling the custom gradient.
  
                  # Expand into (..., 1, d), so that we can matmul it with Jacobian.
                  ws = [tf.expand_dims(w, axis=-2) for w in ws]
                  ws = [ws]  # expand dims on block level as well.
  
                  js, jp = jac
                  ws_js, ws_jp = _block_matmul(ws, js), _block_matmul(
                      ws, jp
                  )
  
                  # Now undo the expansions
                  ws_js, ws_jp = ws_js[0], ws_jp[0]
                  ws_js = [tf.squeeze(t, axis=-2) for t in ws_js]
                  # These should be 0-dimensional
                  ws_jp = [tf.reduce_sum(t) for t in ws_jp]
  
                  # Flatten into a single tuple, so that it has the same structure as args
                  # in inner().
                  return ws_js + ws_jp
  
              return state, gradient
  
          # tf.custom_gradient can only handle a flat sequence of args.
          args = tuple(initial_state + params)
          return inner(*args)
  
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 1 for calling a global function with example usage: for_loop(body_fn, initial_state, params, num_iterations, name) and returns: _accumulating_for_loop, inner, state,  gradient, unknown, unknown,  next_state,  next_jac, unknown