def reduce_max(x, axis, keepdims):
    """Reduces input_tensor along the dimensions given in axis
    Parameters
    ----------
    x: tensor to reduce
    axis: dimensions to reduce, python list
    keepdims: if true, retains reduced dimensions with length 1

    Returns
    -------
    x_red: reduced tensor
    """
    x_red = x
    for n in axis:
        x_red = x_red.max(dim=n, keepdim=keepdims).values
    return x_red