Method test_upsample

python/tests/test_nn.py:1193–1413

(self)

        self.assertEqual(y.dtype, mx.float16)

    def test_upsample(self):
        b, h, w, c = 1, 2, 2, 1
        scale_factor = 2
        upsample_nearest = nn.Upsample(
            scale_factor=scale_factor, mode="nearest", align_corners=True
        )
        upsample_bilinear = nn.Upsample(
            scale_factor=scale_factor, mode="linear", align_corners=True
        )
        upsample_nearest = nn.Upsample(
            scale_factor=scale_factor, mode="nearest", align_corners=True
        )
        upsample_bilinear_no_align_corners = nn.Upsample(
            scale_factor=scale_factor, mode="linear", align_corners=False
        )
        upsample_nearest_no_align_corners = nn.Upsample(
            scale_factor=scale_factor, mode="nearest", align_corners=False
        )
        # Test single feature map, align corners
        x = mx.arange(b * h * w * c).reshape((b, c, h, w)).transpose((0, 2, 3, 1))
        expected_nearest = mx.array(
            [[[[0, 0, 1, 1], [0, 0, 1, 1], [2, 2, 3, 3], [2, 2, 3, 3]]]]
        ).transpose((0, 2, 3, 1))
        expected_bilinear = mx.array(
            [
                [
                    [
                        [0, 0.333333, 0.666667, 1],
                        [0.666667, 1, 1.33333, 1.66667],
                        [1.33333, 1.66667, 2, 2.33333],
                        [2, 2.33333, 2.66667, 3],
                    ]
                ]
            ]
        ).transpose((0, 2, 3, 1))
        # Test single feature map, no align corners
        x = (
            mx.arange(1, b * h * w * c + 1)
            .reshape((b, c, h, w))
            .transpose((0, 2, 3, 1))
        )
        expected_bilinear_no_align_corners = mx.array(
            [
                [
                    [
                        [1.0000, 1.2500, 1.7500, 2.0000],
                        [1.5000, 1.7500, 2.2500, 2.5000],
                        [2.5000, 2.7500, 3.2500, 3.5000],
                        [3.0000, 3.2500, 3.7500, 4.0000],
                    ]
                ]
            ]
        ).transpose((0, 2, 3, 1))
        expected_nearest_no_align_corners = mx.array(
            [[[[1, 1, 2, 2], [1, 1, 2, 2], [3, 3, 4, 4], [3, 3, 4, 4]]]]
        ).transpose((0, 2, 3, 1))
        self.assertTrue(
            np.allclose(