# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import torch

from neuralcompression.functional import soft_round, soft_round_inverse


def test_soft_round_inverse():
    x = torch.linspace(-2.0, 2.0, 50)

    torch.testing.assert_close(
        x,
        soft_round_inverse(x, alpha=1e-13),
    )

    x = torch.tensor([-1.25, -0.75, 0.75, 1.25])

    torch.testing.assert_close(
        x,
        soft_round_inverse(soft_round(x, alpha=2.0), alpha=2.0),
    )

    for offset in range(-5, 5):
        x = torch.linspace(offset + 0.001, offset + 0.999, 100)

        torch.testing.assert_close(
            torch.ceil(x) - 0.5,
            soft_round_inverse(x, alpha=5000.0),
            atol=0.001,
            rtol=0.002,
        )

[((12, 8, 12, 37), 'torch.linspace', 'torch.linspace', ({(12, 23, 12, 27): '-2.0', (12, 29, 12, 32): '2.0', (12, 34, 12, 36): '50'}, {}), '(-2.0, 2.0, 50)', False, 'import torch\n'), ((19, 8, 19, 48), 'torch.tensor', 'torch.tensor', ({(19, 21, 19, 47): '[-1.25, -0.75, 0.75, 1.25]'}, {}), '([-1.25, -0.75, 0.75, 1.25])', False, 'import torch\n'), ((16, 8, 16, 42), 'neuralcompression.functional.soft_round_inverse', 'soft_round_inverse', (), '', False, 'from neuralcompression.functional import soft_round, soft_round_inverse\n'), ((27, 12, 27, 63), 'torch.linspace', 'torch.linspace', ({(27, 27, 27, 41): 'offset + 0.001', (27, 43, 27, 57): 'offset + 0.999', (27, 59, 27, 62): '100'}, {}), '(offset + 0.001, offset + 0.999, 100)', False, 'import torch\n'), ((23, 27, 23, 51), 'neuralcompression.functional.soft_round', 'soft_round', (), '', False, 'from neuralcompression.functional import soft_round, soft_round_inverse\n'), ((31, 12, 31, 47), 'neuralcompression.functional.soft_round_inverse', 'soft_round_inverse', (), '', False, 'from neuralcompression.functional import soft_round, soft_round_inverse\n'), ((30, 12, 30, 25), 'torch.ceil', 'torch.ceil', ({(30, 23, 30, 24): 'x'}, {}), '(x)', False, 'import torch\n')]