code: def _accumulating_for_loop(
      body_fn, initial_state, params, num_iterations, name=None
  ):
      """Version of for_loop with multiple values of num_iterations."""
      # Every tensor in nested tensors (state and Jacobian) gets an extra
      # "accumulating" dimension in front. Functions _create_accumulators etc. below
      # help to work with this dimension.
  
      with tf.name_scope(name or "accumulating_for_loop"):
          max_iterations = tf.math.reduce_max(num_iterations)
          acc_size = tff_utils.get_shape(num_iterations)[0]
  
          # num_iteration = [2, 5] -> mask = [0, 0, 1, 0, 0, 1]. Tells when we should
          # increment acc index before writing. Last element won't be used (i = 0..4).
          mask = tf.scatter_nd(
              indices=tf.expand_dims(num_iterations, axis=-1),
              updates=tf.ones_like(num_iterations),
              shape=(max_iterations + 1,),
          )
          n = len(initial_state)
  
          @tf.custom_gradient
          def inner(*args):
              initial_state, params = args[:n], args[n:]
  
              def while_cond(
                  i, acc_index, state, jac, acc_state, acc_jac
              ):
                  del acc_index, state, jac, acc_state, acc_jac
                  return i < max_iterations
  
              def while_body(
                  i, acc_index, state, jac, acc_state, acc_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)
                  acc_index += mask[i]
                  acc_state = _write_to_accumulators(
                      acc_state, next_state, acc_index
                  )
                  acc_jac = _write_to_accumulators(
                      acc_jac, next_jac, acc_index
                  )
  
                  return (
                      i + 1,
                      acc_index,
                      next_state,
                      next_jac,
                      acc_state,
                      acc_jac,
                  )
  
              initial_acc_state = _create_accumulators(
                  initial_state, acc_size
              )
              initial_acc_state = _write_to_accumulators(
                  initial_acc_state, initial_state, 0
              )
  
              initial_jac = _make_unit_jacobian(initial_state, params)
              initial_acc_jac = _create_accumulators(
                  initial_jac, acc_size
              )
              initial_acc_jac = _write_to_accumulators(
                  initial_acc_jac, initial_jac, 0
              )
  
              loop_vars = (
                  0,
                  0,
                  initial_state,
                  initial_jac,
                  initial_acc_state,
                  initial_acc_jac,
              )
  
              (
                  _,
                  _,
                  _,
                  _,
                  final_acc_state,
                  final_acc_jac,
              ) = tf.compat.v2.while_loop(
                  while_cond,
                  while_body,
                  loop_vars=loop_vars,
                  maximum_iterations=max_iterations,
              )
              # Stack `tf.TensorArray`s into a `Tensor`
              final_acc_state = _stack_accumulators(final_acc_state)
              final_acc_jac = _stack_accumulators(final_acc_jac)
  
              def gradient(*ws):
                  # Same as in for_loop, except we need to sum over the accumulating
                  # dimension. E.g. if x = for_loop(... num_iterations=[2, 5]) and
                  # y = 2*x[0] + 3*x[1], then taking gradient of y will lead to ws having
                  # coeffs 2 and 3 in the acc dimension, and we should sum over it.
                  ws = [tf.expand_dims(w, axis=-2) for w in ws]
                  ws = [ws]  # expand dims on block level as well.
  
                  js, jp = final_acc_jac
                  ws_js, ws_jp = _block_matmul(ws, js), _block_matmul(
                      ws, jp
                  )
  
                  ws_js, ws_jp = ws_js[0], ws_jp[0]
                  ws_js = [tf.squeeze(t, axis=-2) for t in ws_js]
                  ws_jp = [tf.squeeze(t, axis=[-2, -1]) for t in ws_jp]
  
                  # Sum over acc axis.
                  ws_js = [tf.math.reduce_sum(t, axis=0) for t in ws_js]
                  # ws_jp should be 0-dimensional
                  ws_jp = [tf.math.reduce_sum(t) for t in ws_jp]
  
                  return ws_js + ws_jp
  
              return final_acc_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 10 for calling a global function with example usage: _accumulating_for_loop(body_fn, initial_state, params, num_iterations, name) and returns: inner, final_acc_state,  gradient, unknown, unknown,  acc_index,  next_state,  next_jac,  acc_state,  acc_jac, unknown