def METHOD_NAME(
    x: Tensor,
    *,
    mask_axis: Dim,
    broadcast_axis: Union[Dim, Collection[Dim]],
    min_num: Union[int, Tensor],
    max_num: Union[int, Tensor],
    max_dims: Union[int, Tensor],
    mask_value: Union[int, float, Tensor] = 0.0,
) -> Tensor:
    """
    :param x: (batch,time,feature)
    :param mask_axis: axis to mask
    :param broadcast_axis: one or multiple, which should be broadcasted over.
      The remaining axes not specified by mask_axis and broadcast_axis are not broadcasted over
      and treated as batch dims.
      E.g. in [B,T,D], with mask_axis=F, broadcast_axis=T, it creates masks [B,F].
    :param min_num:
    :param max_num: inclusive
    :param max_dims: inclusive
    :param mask_value:
    """
    batch_dims = list(x.dims)
    batch_dims.remove(mask_axis)
    if isinstance(broadcast_axis, Dim):
        batch_dims.remove(broadcast_axis)
    else:
        for a in broadcast_axis:
            batch_dims.remove(a)
    if isinstance(min_num, int) and isinstance(max_num, int) and min_num == max_num:
        num = min_num
    else:
        num = rf.random_uniform(batch_dims, minval=min_num, maxval=max_num + 1, dtype="int32")
    _, indices, k_dim = rf.top_k(
        rf.random_uniform(batch_dims + [mask_axis], minval=0.0, maxval=1.0),
        axis=mask_axis,
        k=num if isinstance(num, int) else rf.reduce_max(num, axis=num.dims),
    )
    # indices should be sorted, and of shape (batch,num), entries (int32) in [0,dim)
    if isinstance(num, int):
        for i in range(num):
            x = mask(
                x,
                mask_axis=mask_axis,
                pos=rf.gather(indices, axis=k_dim, indices=i),
                max_amount=max_dims,
                mask_value=mask_value,
            )
    else:
        def _body(s):
            i_, x_ = s
            y = mask(
                x_,
                mask_axis=mask_axis,
                pos=rf.gather(indices, axis=k_dim, indices=i_),
                max_amount=max_dims,
                mask_value=mask_value,
            )
            y = rf.where(rf.less(i_, num), y, x_)
            return i_ + 1, y
        _, x = rf.while_loop(
            cond=lambda s: s[0] < rf.reduce_max(num, axis=num.dims),
            body=_body,
            initial=(rf.constant(0, dims=()), x),
        )
    return x