def drop_path(
    input: tf.Tensor, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
) -> tf.Tensor:
    """
    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
    """
    if drop_prob == 0.0 or not training:
        return input
    keep_prob = 1 - drop_prob
    input_shape = shape_list(input)
    ndim = len(input_shape)
    shape = [input_shape[0]] + [1] * (ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = tf.random.uniform(shape)
    random_tensor = tf.where(random_tensor <= keep_prob, 1.0, 0.0)
    if keep_prob > 0.0 and scale_by_keep:
        random_tensor /= keep_prob
    return input * random_tensor

    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
    
dpkd/transformers/src/transformers/models/swin/modeling_tf_swin.py