mct-nightly 2.2.0.20241106.458__py3-none-any.whl → 2.2.0.20241108.459__py3-none-any.whl

tests_pytest/keras/gptq/test_gradual_act_quantization.py

@@ -1,102 +0,0 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-from unittest.mock import Mock
-import pytest
-import numpy as np
-import tensorflow as tf
-
-from model_compression_toolkit.gptq.common.gradual_activation_quantization import GradualActivationQuantizerWrapper, \
-    get_gradual_activation_quantizer_wrapper_factory
-from model_compression_toolkit.trainable_infrastructure.keras.annealing_schedulers import KerasLinearAnnealingScheduler
-from model_compression_toolkit.gptq import GradientPTQConfig, GradualActivationQuantizationConfig, QFractionLinearAnnealingConfig
-
-
-
-@pytest.fixture
-def x():
-    return tf.random.normal((2, 5, 6, 7), seed=42, dtype=tf.float32)
-
-
-class Quantizer:
-    def __call__(self, x, training):
-        self.training = training
-        return 3 * x + 1
-
-
-class TestGradualActivationQuantization:
-
-    def test_gradual_act_quant_wrapper(self, x):
-        quantizer = Quantizer()
-        qw = GradualActivationQuantizerWrapper(quantizer, q_fraction_scheduler=lambda t: t / (t + 1))
-
-        y0, y1, y2 = [qw(x, training=True) for _ in range(3)]
-        assert np.allclose(y0.numpy(), x.numpy())  # t=0
-        assert np.allclose(y1.numpy(), 0.5 * x.numpy() + (1.5 * x.numpy() + 0.5))  # t=1
-        assert np.allclose(y2.numpy(), x.numpy() / 3 + (2 * x.numpy() + 2 / 3)) # t=2
-        assert quantizer.training is True
-
-        _ = qw(x, training=False)
-        assert quantizer.training is False  # correct flag was propagated
-
-    def test_factory_no_qdrop(self):
-        quantizer_wrapper, quantizer = self._run_factory_test(qdrop_cfg=None, get_grad_steps_fn=None)
-        assert quantizer_wrapper is quantizer
-
-    @pytest.mark.parametrize('end_step', (20, None))
-    def test_factory_linear(self, x, end_step):
-        qdrop_cfg = GradualActivationQuantizationConfig(
-            QFractionLinearAnnealingConfig(initial_q_fraction=0.3, target_q_fraction=0.8, start_step=10,
-                                           end_step=end_step)
-        )
-
-        def get_total_steps():
-            if end_step is None:
-                return 50
-            assert False  # should not be called if end_step is passed
-
-        quantizer_wrapper, quantizer = self._run_factory_test(qdrop_cfg, get_total_steps)
-
-        scheduler = quantizer_wrapper.q_fraction_scheduler
-        assert isinstance(scheduler, KerasLinearAnnealingScheduler)
-        exp_end_step = 50 if end_step is None else end_step
-        assert scheduler.t_start == 10
-        assert scheduler.t_end == exp_end_step
-        assert scheduler.initial_val == 0.3
-        assert scheduler.target_val == 0.8
-
-        y = [quantizer_wrapper(x, training=True) for _ in range(exp_end_step + 1)]
-
-        assert np.allclose(y[9].numpy(), 0.7 * x.numpy() + 0.3 * quantizer(x, training=True).numpy())
-        assert np.allclose(y[10].numpy(), 0.7 * x.numpy() + 0.3 * quantizer(x, training=True).numpy())
-        assert np.allclose(y[-1].numpy(), 0.2 * x.numpy() + 0.8 * quantizer(x, training=True).numpy())
-
-    def test_factory_linear_common_case(self, x):
-        # validate that we actually implemented the right thing - on first call float input, on last call fully quantized
-        qdrop_cfg = GradualActivationQuantizationConfig(
-            QFractionLinearAnnealingConfig(initial_q_fraction=0, target_q_fraction=1, start_step=0, end_step=None)
-        )
-        quantizer_wrapper, quantizer = self._run_factory_test(qdrop_cfg, lambda: 15)
-        y0, *_, y_last = [quantizer_wrapper(x, training=True) for _ in range(16)]
-        assert np.array_equal(y0.numpy(), x.numpy())
-        assert np.allclose(y_last.numpy(), quantizer(x, training=True).numpy())
-
-    def _run_factory_test(self, qdrop_cfg, get_grad_steps_fn):
-        # Mocks are used to just pass anything
-        gptq_cfg = GradientPTQConfig(n_epochs=5, optimizer=Mock(), loss=Mock(),
-                                     gradual_activation_quantization_config=qdrop_cfg)
-        factory = get_gradual_activation_quantizer_wrapper_factory(gptq_cfg, get_grad_steps_fn, KerasLinearAnnealingScheduler)
-        quantizer = Quantizer()
-        quantizer_wrapper = factory(quantizer)
-        return quantizer_wrapper, quantizer

tests_pytest/keras/trainable_infrastructure/__init__.py

@@ -1,16 +0,0 @@
-#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-#  Licensed under the Apache License, Version 2.0 (the "License");
-#  you may not use this file except in compliance with the License.
-#  You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-#  Unless required by applicable law or agreed to in writing, software
-#  distributed under the License is distributed on an "AS IS" BASIS,
-#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#  See the License for the specific language governing permissions and
-#  limitations under the License.
-#  ==============================================================================
-
-

tests_pytest/keras/trainable_infrastructure/test_linear_annealing.py

@@ -1,49 +0,0 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-import numpy as np
-import pytest
-
-from model_compression_toolkit.trainable_infrastructure.keras.annealing_schedulers import KerasLinearAnnealingScheduler
-
-
-def test_linear_annealing():
-    scheduler = KerasLinearAnnealingScheduler(t_start=10, t_end=35, initial_val=3.4, target_val=-1.6)
-    for t in [0, 9, 10]:
-        assert _isclose(scheduler(t), 3.4)
-
-    for t in [35, 36, 1000]:
-        assert _isclose(scheduler(t), -1.6)
-
-    assert _isclose(scheduler(11), 3.2)
-    assert _isclose(scheduler(27), 0.)
-    assert _isclose(scheduler(34), -1.4)
-
-
-def test_linear_annealing_ascending():
-    scheduler = KerasLinearAnnealingScheduler(t_start=0, t_end=5, initial_val=-0.5, target_val=1.5)
-    assert _isclose(scheduler(0), -0.5)
-    assert _isclose(scheduler(1), -0.1)
-    assert _isclose(scheduler(4), 1.1)
-    assert _isclose(scheduler(5), 1.5)
-
-
-@pytest.mark.parametrize('start', [5, -1])
-def test_invalid(start):
-    with pytest.raises(ValueError):
-        KerasLinearAnnealingScheduler(t_start=start, t_end=4, initial_val=1, target_val=0)
-
-
-def _isclose(x, y):
-    return np.isclose(x, y)

tests_pytest/pytorch/__init__.py

@@ -1,14 +0,0 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================

tests_pytest/pytorch/core/__init__.py

@@ -1,14 +0,0 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================

tests_pytest/pytorch/core/test_data_util.py

@@ -1,125 +0,0 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-import pytest
-import torch
-import numpy as np
-from torch.utils.data import IterableDataset, Dataset
-
-from model_compression_toolkit.core.pytorch.data_util import (data_gen_to_dataloader, IterableDatasetFromGenerator,
-                                                              FixedDatasetFromGenerator, FixedSampleInfoDataset)
-
-
-@pytest.fixture(scope='session')
-def fixed_dataset():
-    # generate 320 images with data1[i] = i and data2[i] = i+10
-    data1 = np.stack([np.full((3, 30, 20), v) for v in range(320)], axis=0)
-    data2 = np.stack([np.full((10,), v + 10) for v in range(320)], axis=0)
-    return data1, data2
-
-
-@pytest.fixture
-def fixed_gen(fixed_dataset):
-    def f():
-        for i in range(10):
-            yield [fixed_dataset[0][32 * i: 32 * (i + 1)], fixed_dataset[1][32 * i: 32 * (i + 1)]]
-
-    return f
-
-
-def get_random_data_gen_fn(seed=42):
-    """ get gen factory for reproducible gen yielding different samples in each epoch """
-    rng = np.random.default_rng(seed)
-
-    def f():
-        for i in range(10):
-            yield [rng.random((32, 3, 20, 30)), rng.random((32, 10))]
-    return f
-
-
-class TestDataUtil:
-    create_dataloader_fn = data_gen_to_dataloader
-
-    def test_iterable_dataset_from_fixed_gen(self, fixed_gen):
-        """ tests iterable dataset from fixed gen - same samples are generated in each epoch in the same order """
-        ds = IterableDatasetFromGenerator(fixed_gen)
-        assert isinstance(ds, IterableDataset)
-        self._validate_ds_from_fixed_gen(ds, 320)
-
-    def test_iterable_dataset_from_random_gen(self):
-        """ test that dataset samples over epochs are identical to the original data generator """
-        ds = IterableDatasetFromGenerator(get_random_data_gen_fn())
-        pass1 = torch.stack([t[0] for t in ds], dim=0)
-        pass2 = torch.stack([t[0] for t in ds], dim=0)
-
-        gen_fn = get_random_data_gen_fn()
-        # one invocation is used for validation and batch size in dataset, so promote the reference gen for comparison
-        next(gen_fn())
-        gen_pass1 = np.concatenate([t[0] for t in gen_fn()], axis=0)
-        gen_pass2 = np.concatenate([t[0] for t in gen_fn()], axis=0)
-        # check that each pass is identical to corresponding pass in the original gen
-        assert np.allclose(pass1.cpu().numpy(), gen_pass1)
-        assert np.allclose(pass2.cpu().numpy(), gen_pass2)
-        assert not torch.equal(pass1, pass2)
-
-    def test_fixed_dataset_from_fixed_gen_full(self, fixed_gen):
-        ds = FixedDatasetFromGenerator(fixed_gen)
-        assert isinstance(ds, Dataset) and not isinstance(ds, IterableDataset)
-        self._validate_ds_from_fixed_gen(ds, 320)
-
-    def test_fixed_dataset_from_const_gen_subset(self, fixed_gen):
-        ds = FixedDatasetFromGenerator(fixed_gen, n_samples=25)
-        self._validate_ds_from_fixed_gen(ds, 25)
-
-    def test_fixed_dataset_from_random_gen_full(self):
-        ds = FixedDatasetFromGenerator(get_random_data_gen_fn())
-        self._validate_fixed_ds(ds, exp_len=320, exp_batch_size=32)
-
-    def test_fixed_dataset_from_random_gen_subset(self):
-        ds = FixedDatasetFromGenerator(get_random_data_gen_fn(), n_samples=123)
-        self._validate_fixed_ds(ds, exp_len=123, exp_batch_size=32)
-
-    def test_not_enough_samples_in_datagen(self):
-        def gen():
-            yield [np.ones((10, 3))]
-        with pytest.raises(ValueError, match='Not enough samples in the data generator'):
-            FixedDatasetFromGenerator(gen, n_samples=11)
-
-    def test_extra_info_mismatch(self, fixed_gen):
-        with pytest.raises(ValueError, match='Mismatch in the number of samples between samples and complementary data'):
-            FixedSampleInfoDataset([1]*10, [2]*10, [3]*11)
-
-    @pytest.mark.parametrize('ds_cls', [FixedDatasetFromGenerator, IterableDatasetFromGenerator])
-    def test_invalid_gen(self, ds_cls):
-        def gen():
-            yield np.ones((10, 3))
-        with pytest.raises(TypeError, match='Data generator is expected to yield a list of tensors'):
-            ds_cls(gen)
-
-    def _validate_ds_from_fixed_gen(self, ds, exp_len):
-        for _ in range(2):
-            for i, sample in enumerate(ds):
-                assert np.array_equal(sample[0].cpu().numpy(), np.full((3, 30, 20), i))
-                assert np.array_equal(sample[1].cpu().numpy(), np.full((10,), i + 10))
-            assert i == exp_len - 1
-            assert ds.orig_batch_size == 32
-            assert len(ds) == exp_len
-
-    def _validate_fixed_ds(self, ds, exp_len, exp_batch_size):
-        assert isinstance(ds, torch.utils.data.Dataset) and not isinstance(ds, torch.utils.data.IterableDataset)
-        full_pass1 = torch.concat([t[0] for t in ds], dim=0)
-        full_pass2 = torch.concat([t[0] for t in ds], dim=0)
-        assert torch.equal(full_pass1, full_pass2)
-        assert len(ds) == exp_len
-        assert ds.orig_batch_size == exp_batch_size

tests_pytest/pytorch/gptq/__init__.py

@@ -1,14 +0,0 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================

tests_pytest/pytorch/gptq/test_annealing_cfg.py

@@ -1,40 +0,0 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-import pytest
-
-from model_compression_toolkit.gptq import QFractionLinearAnnealingConfig
-
-
-def test_linear_annealing_cfg_validation():
-    with pytest.raises(ValueError, match='Expected.* target_q_fraction <= 1'):
-        QFractionLinearAnnealingConfig(initial_q_fraction=0.1, target_q_fraction=1.1, start_step=0, end_step=None)
-
-    with pytest.raises(ValueError, match='Expected.* 0 <= initial_q_fraction'):
-        QFractionLinearAnnealingConfig(initial_q_fraction=-0.1, target_q_fraction=-0.9, start_step=0, end_step=100)
-
-    with pytest.raises(ValueError, match='Expected.* initial_q_fraction < target_q_fraction'):
-        QFractionLinearAnnealingConfig(initial_q_fraction=0.1, target_q_fraction=0.1, start_step=0, end_step=100)
-
-    with pytest.raises(ValueError, match='Expected.* initial_q_fraction < target_q_fraction'):
-        QFractionLinearAnnealingConfig(initial_q_fraction=0.2, target_q_fraction=0.1, start_step=0, end_step=100)
-
-    with pytest.raises(ValueError, match='Expected.* start_step >= 0'):
-        QFractionLinearAnnealingConfig(initial_q_fraction=0, target_q_fraction=1, start_step=-1, end_step=100)
-
-    with pytest.raises(ValueError, match='Expected.* start_step < end_step'):
-        QFractionLinearAnnealingConfig(initial_q_fraction=0, target_q_fraction=1, start_step=100, end_step=100)
-
-    with pytest.raises(ValueError, match='Expected.* start_step < end_step'):
-        QFractionLinearAnnealingConfig(initial_q_fraction=0, target_q_fraction=1, start_step=100, end_step=99)

tests_pytest/pytorch/gptq/test_gradual_act_quantization.py

@@ -1,100 +0,0 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-from unittest.mock import Mock
-
-import pytest
-import torch
-
-from model_compression_toolkit.core.pytorch.pytorch_device_config import get_working_device
-from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import PytorchLinearAnnealingScheduler
-from model_compression_toolkit.gptq import GradientPTQConfig, GradualActivationQuantizationConfig, QFractionLinearAnnealingConfig
-from model_compression_toolkit.gptq.common.gradual_activation_quantization import (
-    GradualActivationQuantizerWrapper, get_gradual_activation_quantizer_wrapper_factory)
-
-
-@pytest.fixture
-def x():
-    return torch.randn((2, 5, 6, 7), generator=torch.Generator().manual_seed(42)).to(device=get_working_device())
-
-
-class Quantizer:
-    def __call__(self, x, training):
-        self.training = training
-        return 3*x + 1
-
-
-class TestGradualActivationQuantization:
-
-    def test_gradual_act_quant_wrapper(self, x):
-        quantizer = Quantizer()
-        qw = GradualActivationQuantizerWrapper(quantizer, q_fraction_scheduler=lambda t: t / (t + 1))
-
-        y0, y1, y2 = [qw(x) for _ in range(3)]
-        assert torch.equal(y0, x)  # t=0
-        assert torch.allclose(y1, 0.5 * x + (1.5 * x + 0.5))  # t=1
-        assert torch.allclose(y2, x / 3 + (2 * x + 2 / 3))  # t=2
-        assert quantizer.training is True
-
-        _ = qw(x, False)
-        assert quantizer.training is False  # correct flag was propagated
-
-    def test_factory_no_qdrop(self):
-        quantizer_wrapper, quantizer = self._run_factory_test(qdrop_cfg=None, get_grad_steps_fn=None)
-        assert quantizer_wrapper is quantizer
-
-    @pytest.mark.parametrize('end_step', (20, None))
-    def test_factory_linear(self, x, end_step):
-        qdrop_cfg = GradualActivationQuantizationConfig(
-            QFractionLinearAnnealingConfig(initial_q_fraction=0.3, target_q_fraction=0.8, start_step=10, end_step=end_step)
-        )
-
-        def get_total_steps():
-            if end_step is None:
-                return 50
-            assert False  # should not be called if end_step is passed
-
-        quantizer_wrapper, quantizer = self._run_factory_test(qdrop_cfg, get_total_steps)
-
-        scheduler = quantizer_wrapper.q_fraction_scheduler
-        assert isinstance(scheduler, PytorchLinearAnnealingScheduler)
-        exp_end_step = 50 if end_step is None else end_step
-        assert scheduler.t_start == 10
-        assert scheduler.t_end == exp_end_step
-        assert scheduler.initial_val == 0.3
-        assert scheduler.target_val == 0.8
-
-        y = [quantizer_wrapper(x) for _ in range(exp_end_step+1)]
-        assert torch.allclose(y[9], 0.7 * x + 0.3 * quantizer(x, True))
-        assert torch.allclose(y[10], 0.7 * x + 0.3 * quantizer(x, True))
-        assert torch.allclose(y[-1], 0.2 * x + 0.8 * quantizer(x, True))
-
-    def test_factory_linear_common_case(self, x):
-        # validate that we actually implemented the right thing - on first call float input, on last call fully quantized
-        qdrop_cfg = GradualActivationQuantizationConfig(
-            QFractionLinearAnnealingConfig(initial_q_fraction=0, target_q_fraction=1, start_step=0, end_step=None)
-        )
-        quantizer_wrapper, quantizer = self._run_factory_test(qdrop_cfg, lambda: 15)
-        y0, *_, y_last = [quantizer_wrapper(x) for _ in range(16)]
-        assert torch.equal(y0, x)
-        assert torch.allclose(y_last, quantizer(x, True))
-
-    def _run_factory_test(self, qdrop_cfg, get_grad_steps_fn):
-        # Mocks are used to just pass anything
-        gptq_cfg = GradientPTQConfig(n_epochs=5, optimizer=Mock(), loss=Mock(),
-                                     gradual_activation_quantization_config=qdrop_cfg)
-        factory = get_gradual_activation_quantizer_wrapper_factory(gptq_cfg, get_grad_steps_fn, PytorchLinearAnnealingScheduler)
-        quantizer = Quantizer()
-        quantizer_wrapper = factory(quantizer)
-        return quantizer_wrapper, quantizer

tests_pytest/pytorch/trainable_infrastructure/__init__.py

@@ -1,14 +0,0 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================

tests_pytest/pytorch/trainable_infrastructure/test_linear_annealing.py

@@ -1,49 +0,0 @@
-# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-import torch
-import pytest
-
-from model_compression_toolkit.trainable_infrastructure.pytorch.annealing_schedulers import PytorchLinearAnnealingScheduler
-
-
-def test_linear_annealing():
-    scheduler = PytorchLinearAnnealingScheduler(t_start=10, t_end=35, initial_val=3.4, target_val=-1.6)
-    for t in [0, 9, 10]:
-        assert _isclose(scheduler(t), 3.4)
-
-    for t in [35, 36, 1000]:
-        assert _isclose(scheduler(t), -1.6)
-
-    assert _isclose(scheduler(11), 3.2)
-    assert _isclose(scheduler(27), 0.)
-    assert _isclose(scheduler(34), -1.4)
-
-
-def test_linear_annealing_ascending():
-    scheduler = PytorchLinearAnnealingScheduler(t_start=0, t_end=5, initial_val=-0.5, target_val=1.5)
-    assert _isclose(scheduler(0), -0.5)
-    assert _isclose(scheduler(1), -0.1)
-    assert _isclose(scheduler(4), 1.1)
-    assert _isclose(scheduler(5), 1.5)
-
-
-@pytest.mark.parametrize('start', [5, -1])
-def test_invalid(start):
-    with pytest.raises(ValueError):
-        PytorchLinearAnnealingScheduler(t_start=start, t_end=4, initial_val=1, target_val=0)
-
-
-def _isclose(x, y):
-    return torch.isclose(x, torch.tensor(y))