def insert_zeros(tensor: tf.Tensor, mask: tf.Tensor) -> tf.Tensor:
    """Inserts zeros into tensor before each masked element.
    For example:
    ```python
        output = insert_zeros(
            tensor =  tf.constant([[1, 2, 3, 4, 5], [10, 20, 30, 40, 50]], dtype = tf.int32),
            mask = tf.constant([[False, True, False, False, True], [False, True,  True, True,  False]]),
        )
        # -> [[1, 0, 2, 3, 4, 0, 5, 0], [10, 0, 20, 0, 30, 0, 40, 50]]
        # We insert 0s 2, 5, 20, 30, and 40 because their positions in input tensor corresponds to True value
        in mask.
    ```

    Args:
        tensor: tf.Tensor, shape = [batch, length], any type and the same shape as mask
        mask:   tf.Tensor, shape = [batch, length], dtype = tf.bool and the same shape as tensor

    Returns:    tf.Tensor, shape = [batch, length + max_num_insertions],
                where max_num_insertions is the maximal number of True values along the 0 batch dimension of mask.
                dtype = same as input tensor
    """
    batch_size = tf.shape(tensor)[0]
    length = tf.shape(mask)[1]

    delta = tf.cumsum(tf.cast(mask, dtype=tf.int32), exclusive=False, axis=1)
    max_num_insertions = tf.reduce_max(delta[:, -1])

    y, x = tf.meshgrid(tf.range(length), tf.range(batch_size))
    y = y + delta
    indices = tf.reshape(tf.stack([x, y], 2), [-1, 2])

    output = tf.scatter_nd(indices=indices, updates=tf.reshape(tensor, shape=[-1]), shape=tf.stack([batch_size, length + max_num_insertions]))

    return output
Inserts zeros into tensor before each masked element.
    For example:
    ```python
        output = insert_zeros(
            tensor =  tf.constant([[1, 2, 3, 4, 5], [10, 20, 30, 40, 50]], dtype = tf.int32),
            mask = tf.constant([[False, True, False, False, True], [False, True,  True, True,  False]]),
        )
        # -> [[1, 0, 2, 3, 4, 0, 5, 0], [10, 0, 20, 0, 30, 0, 40, 50]]
        # We insert 0s 2, 5, 20, 30, and 40 because their positions in input tensor corresponds to True value
        in mask.
    ```

    Args:
        tensor: tf.Tensor, shape = [batch, length], any type and the same shape as mask
        mask:   tf.Tensor, shape = [batch, length], dtype = tf.bool and the same shape as tensor

    Returns:    tf.Tensor, shape = [batch, length + max_num_insertions],
                where max_num_insertions is the maximal number of True values along the 0 batch dimension of mask.
                dtype = same as input tensor
    
tensorflow_asr/losses/impl/ctc_tpu.py