def unfold(
    init_tensor: tf.Tensor,
    iterfunc: Callable[[tf.Tensor, tf.Tensor], tf.Tensor],
    num_iters: Union[int, tf.Tensor],
    d_i: int,
    element_shape: tf.TensorShape,
    swap_memory: bool = False,
    name: str = "unfold",
) -> tf.Tensor:
    """Calculates a tensor by iterations over i that is the concatenation
        for d_i = +1:
            init_tensor
            iterfunc(init_tensor, 0)
            iterfunc(iterfunc(init_tensor, 0), 1)
            ...
            ..., num_iters - 1)
            ..., num_iters - 1), num_iters)
        for d_i = -1:
            ..., 2), 1), 0)
            ..., 2), 1)
            ...
            iterfunc(iterfunc(init_tensor, num_iters - 1), num_iters - 2)
            iterfunc(init_tensor, num_iters - 1)
            init_tensor
    For example:
    ```python
        unfold(
            init_tensor=tf.constant(0),
            iterfunc=lambda x, i: x + i,
            num_iters=5,
            d_i=1,
            element_shape=tf.TensorShape([]),
        )
        # -> [0, 0, 1, 3, 6, 10]
    ```

    Args:
        init_tensor:    tf.Tensor, of any shape that is the initial value of the iterations.
        iterfunc:       tf.Tensor, tf.Tensor -> tf.Tensor, that is the iteration function
                            from and onto the same shape as init_tensor
        num_iters:      tf.Tensor or static integer that is the number of iterations
        d_i:            either +1 or -1, where
                            +1 corresponds for the iterations from 0 to num_iters inclusive
                            -1 corresponds for the iterations from num_iters to 0 inclusive
        element_shape:  tf.TensorShape([]) that is the shape of init_tensor
        swap_memory:    the same as for tf.while_loop, argument
        name:           str, local tensor names scope

    Returns:            tf.Tensor, shape = [num_iters + 1] + init_tensor.shape
                        dtype the same as init_tensor
    """
    assert d_i in {-1, 1}
    positive_direction = d_i == 1

    with tf.name_scope(name):
        num_iters = tf.convert_to_tensor(num_iters)

        tensor_array = tf.TensorArray(
            dtype=init_tensor.dtype,
            size=num_iters + 1,
            element_shape=element_shape,
            clear_after_read=False,
            infer_shape=True,
            dynamic_size=False,
        )
        tensor_array = tensor_array.write(0 if positive_direction else num_iters, init_tensor)

        def body(i, tensor_slice):
            last_value = tensor_slice.read(i if positive_direction else i + 1)
            new_value = iterfunc(last_value, i)
            tensor_slice = tensor_slice.write(i + 1 if positive_direction else i, new_value)
            return i + d_i, tensor_slice

        n = tf.constant(0, dtype=tf.int32) if positive_direction else num_iters - 1
        _, array_out = tf.while_loop(
            cond=lambda i, _: tf.constant(True),
            body=body,
            loop_vars=(n, tensor_array),
            maximum_iterations=num_iters,
            swap_memory=swap_memory,
            name="unfold_while_loop",
        )
        return array_out.stack()
Calculates a tensor by iterations over i that is the concatenation
        for d_i = +1:
            init_tensor
            iterfunc(init_tensor, 0)
            iterfunc(iterfunc(init_tensor, 0), 1)
            ...
            ..., num_iters - 1)
            ..., num_iters - 1), num_iters)
        for d_i = -1:
            ..., 2), 1), 0)
            ..., 2), 1)
            ...
            iterfunc(iterfunc(init_tensor, num_iters - 1), num_iters - 2)
            iterfunc(init_tensor, num_iters - 1)
            init_tensor
    For example:
    ```python
        unfold(
            init_tensor=tf.constant(0),
            iterfunc=lambda x, i: x + i,
            num_iters=5,
            d_i=1,
            element_shape=tf.TensorShape([]),
        )
        # -> [0, 0, 1, 3, 6, 10]
    ```

    Args:
        init_tensor:    tf.Tensor, of any shape that is the initial value of the iterations.
        iterfunc:       tf.Tensor, tf.Tensor -> tf.Tensor, that is the iteration function
                            from and onto the same shape as init_tensor
        num_iters:      tf.Tensor or static integer that is the number of iterations
        d_i:            either +1 or -1, where
                            +1 corresponds for the iterations from 0 to num_iters inclusive
                            -1 corresponds for the iterations from num_iters to 0 inclusive
        element_shape:  tf.TensorShape([]) that is the shape of init_tensor
        swap_memory:    the same as for tf.while_loop, argument
        name:           str, local tensor names scope

    Returns:            tf.Tensor, shape = [num_iters + 1] + init_tensor.shape
                        dtype the same as init_tensor
    
tensorflow_asr/losses/impl/ctc_tpu.py