Repository: microsoft/DeepSpeed
File: deepspeed/profiling/flops_profiler/profiler.py (Python)

Code before:
    # estimation
    return input.numel() * (5 if has_affine else 4), 0


def _upsample_flops_compute(*args, **kwargs):
    input = args[0]
    size = kwargs.get('size', None)
    if size is None and len(args) > 1:
        size = args[1]

    if size is not None:
        if isinstance(size, tuple) or isinstance(size, list):
            return int(_prod(size)), 0
        else:
            return int(size), 0

    scale_factor = kwargs.get('scale_factor', None)
    if scale_factor is None and len(args) > 2:
        scale_factor = args[2]
    assert scale_factor is not None, "either size or scale_factor should be defined"

    flops = input.numel()
    if isinstance(scale_factor, (list, tuple)):
        # see documentation of `F.interpolate`
        # the spatial dims are defined as the last `n-2` dims of the tensor
        assert len(scale_factor) == input.ndim - 2
        flops *= int(_prod(scale_factor))
    else:
        flops *= scale_factor**(input.ndim - 2)
    return flops, 0


def _softmax_flops_compute(input, dim=None, _stacklevel=3, dtype=None):
    return input.numel(), 0


def _embedding_flops_compute(
    input,
    weight,
    padding_idx=None,
    max_norm=None,
    norm_type=2.0,
    scale_grad_by_freq=False,
    sparse=False,
):
    return 0, 0


def _dropout_flops_compute(input, p=0.5, training=True, inplace=False):
    return 0, 0


Review comment (suggestion):
It is not directly related to this fix, but I'm not sure if this code is right.
Torch's document says `scale_factor` is a float or a tuple of float values. It doesn't seem right to convert `_prod(scale_factor)` to int. If we want to return int, we should have `flops *=_prod(scale_factor)` and `return int(flops), 0` (the else branch also seems to return float)



Code after:
    # estimation
    return input.numel() * (5 if has_affine else 4), 0


def _upsample_flops_compute(*args, **kwargs):
    input = args[0]
    size = kwargs.get('size', None)
    if size is None and len(args) > 1:
        size = args[1]

    if size is not None:
        if isinstance(size, tuple) or isinstance(size, list):
            return int(_prod(size)), 0
        else:
            return int(size), 0

    scale_factor = kwargs.get('scale_factor', None)
    if scale_factor is None and len(args) > 2:
        scale_factor = args[2]
    assert scale_factor is not None, "either size or scale_factor should be defined"

    flops = input.numel()
    if isinstance(scale_factor, (list, tuple)):
        # see documentation of `F.interpolate`
        # the spatial dims are defined as the last `n-2` dims of the tensor
        assert len(scale_factor) == input.ndim - 2
        flops *= _prod(scale_factor)
    else:
        flops *= scale_factor**(input.ndim - 2)

    return int(flops), 0


def _softmax_flops_compute(input, dim=None, _stacklevel=3, dtype=None):
    return input.numel(), 0


def _embedding_flops_compute(
    input,
    weight,
    padding_idx=None,
    max_norm=None,
    norm_type=2.0,
    scale_grad_by_freq=False,
    sparse=False,
):
    return 0, 0


def _dropout_flops_compute(input, p=0.5, training=True, inplace=False):
    return 0, 0