def rel_left_shift(x, causal=False):
    """
    Relative left shift

    Input:
    tf.Tensor(
    [[1 2 3]
    [4 5 6]
    [7 8 9]], shape=(3, 3), dtype=int32)

    Output:
    tf.Tensor(
    [[3 0 0]
    [5 6 0]
    [7 8 9]], shape=(3, 3), dtype=int32)

    Args:
        x (tf.Tensor): shape BNTR

    Returns:
        x: left shifted, shape BNTR
    """
    b, n, t, r = shape_util.shape_list(x)

    # fmt: off
    if causal:
        x = tf.pad(x, [[0, 0], [0, 0], [0, 0], [1, 0]]) # [B, N, T, Th + T]
        x = tf.reshape(x, [b, n, -1])
        x = tf.pad(x, [[0, 0], [0, 0], [r - t, 0]])
        x = tf.reshape(x, [b, n, 1 + t, r])
        x = tf.slice(x, begin=[0, 0, 1, 0], size=[-1, -1, -1, -1]) # [B, N, T, Th + T]
    else:
        x = tf.pad(x, [[0, 0], [0, 0], [0, 0], [0, 1]])  # [B, N, T, Th + 2*T] where R = Th + 2*T - 1, S = Th + T
        x = tf.reshape(x, [b, n, -1])  # [B, N, TTh + 2*TT]
        x = tf.pad(x, [[0, 0], [0, 0], [0, r - t]])  # [B, N, TTh + 2*TT + Th + 2*T - 1 - T] = [B, N, TTh + 2*TT + Th + T - 1]
        x = tf.reshape(x, [b, n, 1 + t, r])  # TTh + 2*TT + Th + T - 1 = TTh + 2*TT + Th + 2*T - T - 1 = Th(T + 1) + 2*T(T + 1) - (T + 1) = (T + 1)(Th + 2*T - 1) = (T + 1)R # pylint: disable=line-too-long
        x = tf.slice(x, begin=[0, 0, 0, (t - 1)], size=[-1, -1, t, -1]) # [B, N, T, Th + T]
    # fmt: on

    # x = tf.transpose(x, perm=[2, 3, 0, 1])  # BNTR -> TRBN
    # x_shape = tf.shape(x)

    # x = tf.pad(x, [[0, 0], [1, 0], [0, 0], [0, 0]])  # shift on position time dimension R
    # x = tf.reshape(x, [x_shape[1] + 1, x_shape[0], x_shape[2], x_shape[3]])
    # x = tf.slice(x, [1, 0, 0, 0], [-1, -1, -1, -1])
    # x = tf.reshape(x, x_shape)
    # if mask_upper_triangle:
    #     x *= tf.reverse(tf.linalg.band_part(tf.ones((x_shape[0], x_shape[1]), x.dtype), 0, -1), [0, 1])[..., tf.newaxis, tf.newaxis]

    # x = tf.transpose(x, perm=[2, 3, 0, 1])  # TRBN -> BNTR
    return x

    Relative left shift

    Input:
    tf.Tensor(
    [[1 2 3]
    [4 5 6]
    [7 8 9]], shape=(3, 3), dtype=int32)

    Output:
    tf.Tensor(
    [[3 0 0]
    [5 6 0]
    [7 8 9]], shape=(3, 3), dtype=int32)

    Args:
        x (tf.Tensor): shape BNTR

    Returns:
        x: left shifted, shape BNTR
    
tensorflow_asr/models/layers/multihead_attention.py