| |
|
|
| import math |
| import unittest |
| from itertools import permutations |
|
|
| import mlx.core as mx |
| import mlx_tests |
| import numpy as np |
|
|
| try: |
| import torch |
| import torch.nn.functional as F |
|
|
| has_torch = True |
| except ImportError as e: |
| has_torch = False |
|
|
|
|
| class TestConvTranspose(mlx_tests.MLXTestCase): |
| @unittest.skipIf(not has_torch, "requires Torch") |
| def test_torch_conv_transpose_1D(self): |
| def run_conv_transpose_1D( |
| N, |
| C, |
| O, |
| iH, |
| kH, |
| stride, |
| padding, |
| output_padding=0, |
| dilation=1, |
| groups=1, |
| dtype="float32", |
| atol=1e-5, |
| ): |
| with self.subTest( |
| dtype=dtype, |
| N=N, |
| C=C, |
| O=O, |
| iH=iH, |
| kH=kH, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ): |
| np_dtype = getattr(np, dtype) |
| np.random.seed(0) |
| in_np = np.random.normal(0, 1.0 / C, (N, iH, C)).astype(np_dtype) |
| wt_np = np.random.normal(0, 1.0 / C, (O, kH, int(C / groups))).astype( |
| np_dtype |
| ) |
|
|
| in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| in_pt = torch.from_numpy(in_np.transpose(0, 2, 1)) |
| wt_pt = torch.from_numpy(wt_np.transpose(2, 0, 1)) |
|
|
| out_mx = mx.conv_transpose1d( |
| in_mx, |
| wt_mx, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
| out_pt = torch.conv_transpose1d( |
| in_pt, |
| wt_pt, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
| out_pt = torch.transpose(out_pt, 2, 1) |
|
|
| self.assertEqual(out_pt.shape, out_mx.shape) |
| self.assertTrue(np.allclose(out_pt.numpy(), out_mx, atol=atol)) |
|
|
| for dtype in ("float32",): |
| for N, C, O in ( |
| (1, 1, 1), |
| (1, 6, 1), |
| (1, 1, 6), |
| (4, 32, 64), |
| ): |
| for iH, kH, stride, padding in ( |
| (1, 1, 1, 0), |
| (3, 3, 1, 0), |
| (31, 5, 5, 2), |
| ): |
| run_conv_transpose_1D(N, C, O, iH, kH, stride, padding, dtype=dtype) |
|
|
| |
| N, C, O = (4, 32, 64) |
| for iH, kH, stride, padding in ( |
| (1, 1, 1, 0), |
| (3, 3, 1, 0), |
| (31, 5, 5, 2), |
| ): |
| for group in (1,): |
| run_conv_transpose_1D( |
| N, C, O, iH, kH, stride, padding, groups=group, dtype=dtype |
| ) |
|
|
| |
| for tpose_in, tpose_wt in ( |
| ((0, 2, 1), (0, 1, 2)), |
| ((0, 2, 1), (0, 2, 1)), |
| ): |
| with self.subTest(name="strided", tpose_in=tpose_in, tpose_wt=tpose_wt): |
| in_np = np.random.normal(0, 1.0 / 16, (16, 16, 16)).astype(np.float32) |
| wt_np = np.random.normal(0, 1.0 / 16, (16, 16, 16)).astype(np.float32) |
|
|
| in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| in_mx_t = mx.transpose(in_mx, tpose_in) |
| wt_mx_t = mx.transpose(wt_mx, tpose_wt) |
| out_mx = mx.conv_transpose1d(in_mx_t, wt_mx_t) |
|
|
| in_pt = torch.from_numpy(in_np.transpose(tpose_in).transpose(0, 2, 1)) |
| wt_pt = torch.from_numpy(wt_np.transpose(tpose_wt).transpose(2, 0, 1)) |
|
|
| out_pt = torch.conv_transpose1d(in_pt, wt_pt) |
| out_pt = torch.transpose(out_pt, 2, 1) |
|
|
| self.assertEqual(out_pt.shape, out_mx.shape) |
| self.assertTrue(np.allclose(out_pt.numpy(), out_mx, atol=1e-5)) |
|
|
| @unittest.skipIf(not has_torch, "requires Torch") |
| def test_torch_conv_transpose_1D_grad(self): |
| def run_conv_transpose1D_grad( |
| N, |
| C, |
| O, |
| iH, |
| kH, |
| stride, |
| padding, |
| dilation=1, |
| groups=1, |
| dtype="float32", |
| atol=1e-5, |
| ): |
| with self.subTest( |
| dtype=dtype, |
| N=N, |
| C=C, |
| O=O, |
| iH=iH, |
| kH=kH, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ): |
| np_dtype = getattr(np, dtype) |
| np.random.seed(0) |
| |
|
|
| in_np = np.random.normal(0, 1.0 / C, (N, iH, C)).astype(np_dtype) |
| wt_np = np.random.normal(0, 1.0 / C, (O, kH, C)).astype(np_dtype) |
|
|
| in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| in_pt = torch.from_numpy(in_np.transpose(0, 2, 1)).requires_grad_(True) |
| wt_pt = torch.from_numpy(wt_np.transpose(2, 0, 1)).requires_grad_(True) |
|
|
| out_pt = F.conv_transpose1d( |
| in_pt, wt_pt, stride=stride, padding=padding, dilation=dilation |
| ) |
|
|
| |
| out_pt.retain_grad() |
| out_pt.sum().backward() |
|
|
| pt_grad_in = in_pt.grad.permute(0, 2, 1).numpy() |
| pt_grad_wt = wt_pt.grad.permute(1, 2, 0).numpy() |
|
|
| ct_mx = mx.array(out_pt.grad.numpy().transpose(0, 2, 1)) |
|
|
| def f(a, b): |
| return mx.conv_transpose1d( |
| a, |
| b, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
|
|
| _, outs_mx = mx.vjp( |
| f, |
| [ |
| in_mx, |
| wt_mx, |
| ], |
| [ |
| ct_mx, |
| ], |
| ) |
|
|
| mx_grad_in, mx_grad_wt = outs_mx |
|
|
| self.assertEqual(pt_grad_in.shape, mx_grad_in.shape) |
| self.assertEqual(in_mx.shape, mx_grad_in.shape) |
| self.assertTrue(np.allclose(pt_grad_in, mx_grad_in, atol=atol)) |
|
|
| self.assertEqual(pt_grad_wt.shape, mx_grad_wt.shape) |
| self.assertEqual(wt_mx.shape, mx_grad_wt.shape) |
| self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) |
|
|
| for dtype in ("float32",): |
| for N, C, O in ( |
| (1, 1, 1), |
| (1, 6, 1), |
| (1, 1, 6), |
| (4, 32, 64), |
| ): |
| for iH, kH, stride, padding in ( |
| (1, 1, 1, 0), |
| (3, 3, 1, 0), |
| (31, 5, 5, 2), |
| ): |
| run_conv_transpose1D_grad( |
| N, C, O, iH, kH, stride, padding, dtype=dtype |
| ) |
|
|
| @unittest.skipIf(not has_torch, "requires Torch") |
| def test_torch_conv_transpose_2D(self): |
| def run_conv_transpose2D( |
| N, |
| C, |
| O, |
| idim, |
| kdim, |
| stride, |
| padding, |
| dilation=(1, 1), |
| groups=1, |
| dtype="float32", |
| atol=1e-5, |
| ): |
| with self.subTest( |
| dtype=dtype, |
| N=N, |
| C=C, |
| O=O, |
| idim=idim, |
| kdim=kdim, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ): |
| np_dtype = getattr(np, dtype) |
| np.random.seed(0) |
| iH, iW = idim |
| kH, kW = kdim |
| scale = 1.0 / math.sqrt(kH * kW * C) |
| in_np = np.random.normal(0.0, scale, (N, iH, iW, C)).astype(np_dtype) |
| wt_np = np.random.normal(0.0, 1.0, (O, kH, kW, int(C / groups))).astype( |
| np_dtype |
| ) |
|
|
| in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)).to("cpu") |
| wt_pt = torch.from_numpy(wt_np.transpose(3, 0, 1, 2)).to("cpu") |
|
|
| out_mx = mx.conv_transpose2d( |
| in_mx, |
| wt_mx, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
| out_pt = torch.conv_transpose2d( |
| in_pt, |
| wt_pt, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
| out_pt = torch.permute(out_pt, (0, 2, 3, 1)).numpy(force=True) |
|
|
| self.assertEqual(out_pt.shape, out_mx.shape) |
| self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) |
|
|
| for dtype in ("float32",): |
| for N, C, O in ( |
| (1, 1, 1), |
| (1, 6, 1), |
| (1, 1, 6), |
| (4, 32, 64), |
| ): |
| for idim, kdim, stride, padding in ( |
| ((1, 1), (1, 1), (1, 1), (0, 0)), |
| ((3, 3), (3, 1), (1, 1), (0, 0)), |
| ((31, 31), (5, 5), (5, 5), (2, 2)), |
| ): |
| run_conv_transpose2D( |
| N, C, O, idim, kdim, stride, padding, dtype=dtype |
| ) |
|
|
| |
| N, C, O = (4, 32, 64) |
| for idim, kdim, stride, padding in ( |
| ((1, 1), (1, 1), (1, 1), (0, 0)), |
| ((3, 3), (3, 1), (1, 1), (0, 0)), |
| ((31, 31), (5, 5), (5, 5), (2, 2)), |
| ): |
| for group in (1,): |
| run_conv_transpose2D( |
| N, C, O, idim, kdim, stride, padding, groups=group, dtype=dtype |
| ) |
|
|
| @unittest.skipIf(not has_torch, "requires Torch") |
| def test_torch_conv_transpose_2D_grad(self): |
| def run_conv_transpose2D_grad( |
| N, |
| C, |
| O, |
| idim, |
| kdim, |
| stride, |
| padding, |
| dilation=(1, 1), |
| groups=1, |
| dtype="float32", |
| atol=1e-5, |
| ): |
| with self.subTest( |
| dtype=dtype, |
| N=N, |
| C=C, |
| O=O, |
| idim=idim, |
| kdim=kdim, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ): |
| np_dtype = getattr(np, dtype) |
| np.random.seed(0) |
| iH, iW = idim |
| kH, kW = kdim |
| scale = 1.0 / math.sqrt(kH * kW * C * O) |
|
|
| in_np = np.random.normal(0.0, scale, (N, iH, iW, C)).astype(np_dtype) |
| wt_np = np.random.normal(0.0, scale, (O, kH, kW, C)).astype(np_dtype) |
|
|
| in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)).requires_grad_( |
| True |
| ) |
| wt_pt = torch.from_numpy(wt_np.transpose(3, 0, 1, 2)).requires_grad_( |
| True |
| ) |
|
|
| out_pt = F.conv_transpose2d( |
| in_pt, wt_pt, stride=stride, padding=padding, dilation=dilation |
| ) |
|
|
| |
| out_pt.retain_grad() |
| out_pt.sum().backward() |
|
|
| pt_grad_in = in_pt.grad.permute(0, 2, 3, 1).numpy() |
| pt_grad_wt = wt_pt.grad.permute(1, 2, 3, 0).numpy() |
|
|
| ct_mx = mx.array(out_pt.grad.numpy().transpose(0, 2, 3, 1)) |
|
|
| def f(a, b): |
| return mx.conv_transpose2d( |
| a, |
| b, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
|
|
| _, outs_mx = mx.vjp( |
| f, |
| [in_mx, wt_mx], |
| [ct_mx], |
| ) |
|
|
| mx_grad_in, mx_grad_wt = outs_mx |
|
|
| self.assertEqual(pt_grad_in.shape, mx_grad_in.shape) |
| self.assertEqual(in_mx.shape, mx_grad_in.shape) |
| self.assertTrue(np.allclose(pt_grad_in, mx_grad_in, atol=atol)) |
|
|
| self.assertEqual(pt_grad_wt.shape, mx_grad_wt.shape) |
| self.assertEqual(wt_mx.shape, mx_grad_wt.shape) |
| self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) |
|
|
| for dtype in ("float32",): |
| for N, C, O in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (4, 32, 64), (4, 16, 32)): |
| for idim, kdim, stride, padding, dilation in ( |
| ((1, 1), (1, 1), (1, 1), (0, 0), (1, 1)), |
| ((3, 3), (3, 1), (1, 1), (0, 0), (1, 1)), |
| ((31, 31), (5, 5), (5, 5), (2, 2), (1, 1)), |
| ((32, 32), (3, 3), (2, 2), (1, 1), (1, 1)), |
| ((31, 31), (5, 5), (5, 5), (2, 2), (3, 2)), |
| ((32, 32), (3, 3), (2, 2), (1, 1), (3, 2)), |
| ): |
| run_conv_transpose2D_grad( |
| N, C, O, idim, kdim, stride, padding, dilation, dtype=dtype |
| ) |
|
|
| @unittest.skipIf(not has_torch, "requires Torch") |
| def test_torch_conv_transpose_3D(self): |
| def run_conv_transpose3D( |
| N, |
| C, |
| O, |
| idim, |
| kdim, |
| stride, |
| padding, |
| dilation=(1, 1, 1), |
| groups=1, |
| dtype="float32", |
| atol=1e-5, |
| ): |
| with self.subTest( |
| dtype=dtype, |
| N=N, |
| C=C, |
| O=O, |
| idim=idim, |
| kdim=kdim, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ): |
| np_dtype = getattr(np, dtype) |
| np.random.seed(0) |
| iD, iH, iW = idim |
| kD, kH, kW = kdim |
| scale = 1.0 / math.sqrt(kD * kH * kW * C * O) |
| in_np = np.random.normal(0.0, scale, (N, iD, iH, iW, C)).astype( |
| np_dtype |
| ) |
| wt_np = np.random.normal(0.0, 1.0, (O, kD, kH, kW, C)).astype(np_dtype) |
|
|
| in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| in_pt = torch.from_numpy(in_np.transpose(0, 4, 1, 2, 3)) |
| wt_pt = torch.from_numpy(wt_np.transpose(4, 0, 1, 2, 3)) |
|
|
| out_mx = mx.conv_transpose3d( |
| in_mx, |
| wt_mx, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
| out_pt = torch.conv_transpose3d( |
| in_pt, |
| wt_pt, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
| out_pt = torch.permute(out_pt, (0, 2, 3, 4, 1)).numpy(force=True) |
|
|
| self.assertEqual(out_pt.shape, out_mx.shape) |
| self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) |
|
|
| for dtype in ("float32",): |
| for N, C, O in ( |
| (1, 1, 1), |
| (1, 6, 1), |
| (1, 1, 6), |
| (2, 8, 16), |
| ): |
| for idim, kdim, stride, padding in ( |
| ((1, 1, 1), (1, 1, 1), (1, 1, 1), (0, 0, 0)), |
| ((3, 3, 3), (3, 1, 1), (1, 1, 1), (0, 0, 0)), |
| ((15, 15, 15), (3, 3, 3), (3, 3, 3), (2, 2, 2)), |
| ): |
| run_conv_transpose3D( |
| N, C, O, idim, kdim, stride, padding, dtype=dtype |
| ) |
|
|
| @unittest.skipIf(not has_torch, "requires Torch") |
| def test_torch_conv_transpose_3D_grad(self): |
| def run_conv_transpose3D_grad( |
| N, |
| C, |
| O, |
| idim, |
| kdim, |
| stride, |
| padding, |
| dilation=(1, 1, 1), |
| groups=1, |
| dtype="float32", |
| atol=1e-4, |
| ): |
| with self.subTest( |
| dtype=dtype, |
| N=N, |
| C=C, |
| O=O, |
| idim=idim, |
| kdim=kdim, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ): |
| np_dtype = getattr(np, dtype) |
| np.random.seed(0) |
| iD, iH, iW = idim |
| kD, kH, kW = kdim |
| scale = 1.0 / math.sqrt(kD * kH * kW * C * O) |
|
|
| in_np = np.random.normal(0.0, scale, (N, iD, iH, iW, C)).astype( |
| np_dtype |
| ) |
| wt_np = np.random.normal(0.0, scale, (O, kD, kH, kW, C)).astype( |
| np_dtype |
| ) |
|
|
| in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| in_pt = torch.from_numpy(in_np.transpose(0, 4, 1, 2, 3)).requires_grad_( |
| True |
| ) |
| wt_pt = torch.from_numpy(wt_np.transpose(4, 0, 1, 2, 3)).requires_grad_( |
| True |
| ) |
|
|
| out_pt = F.conv_transpose3d( |
| in_pt, |
| wt_pt, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
|
|
| |
| out_pt.retain_grad() |
| out_pt.sum().backward() |
|
|
| pt_grad_in = in_pt.grad.permute(0, 2, 3, 4, 1).numpy() |
| pt_grad_wt = wt_pt.grad.permute(1, 2, 3, 4, 0).numpy() |
|
|
| ct_mx = mx.array(out_pt.grad.numpy().transpose(0, 2, 3, 4, 1)) |
|
|
| def f(a, b): |
| return mx.conv_transpose3d( |
| a, |
| b, |
| stride=stride, |
| padding=padding, |
| dilation=dilation, |
| groups=groups, |
| ) |
|
|
| _, outs_mx = mx.vjp( |
| f, |
| [in_mx, wt_mx], |
| [ct_mx], |
| ) |
|
|
| mx_grad_in, mx_grad_wt = outs_mx |
|
|
| self.assertEqual(pt_grad_in.shape, mx_grad_in.shape) |
| self.assertEqual(in_mx.shape, mx_grad_in.shape) |
| self.assertTrue(np.allclose(pt_grad_in, mx_grad_in, atol=atol)) |
|
|
| self.assertEqual(pt_grad_wt.shape, mx_grad_wt.shape) |
| self.assertEqual(wt_mx.shape, mx_grad_wt.shape) |
| self.assertTrue(np.allclose(pt_grad_wt, mx_grad_wt, atol=atol)) |
|
|
| for dtype in ("float32",): |
| for N, C, O in ((1, 1, 1), (1, 6, 1), (1, 1, 6), (2, 4, 8), (2, 8, 16)): |
| for idim, kdim, stride, padding, dilation in ( |
| ((1, 1, 1), (1, 1, 1), (1, 1, 1), (0, 0, 0), (1, 1, 1)), |
| ((3, 3, 3), (3, 1, 1), (1, 1, 1), (0, 0, 0), (1, 1, 1)), |
| ((7, 7, 7), (5, 5, 5), (5, 5, 5), (2, 2, 2), (1, 1, 1)), |
| ((8, 8, 8), (3, 3, 3), (2, 2, 2), (1, 1, 1), (1, 1, 1)), |
| ((7, 7, 7), (5, 5, 5), (3, 3, 3), (2, 2, 2), (3, 2, 2)), |
| ((8, 8, 8), (3, 3, 3), (2, 2, 2), (1, 1, 1), (3, 2, 2)), |
| ): |
| run_conv_transpose3D_grad( |
| N, C, O, idim, kdim, stride, padding, dilation, dtype=dtype |
| ) |
|
|
| @unittest.skipIf(not has_torch, "requires Torch") |
| def test_torch_conv_tranpose_1d_output_padding(self): |
| def run_conv_transpose_1d_output_padding( |
| N, C, O, iH, kH, stride, padding, output_padding, dtype="float32", atol=1e-5 |
| ): |
| with self.subTest( |
| dtype=dtype, |
| N=N, |
| C=C, |
| O=O, |
| iH=iH, |
| kH=kH, |
| stride=stride, |
| padding=padding, |
| output_padding=output_padding, |
| ): |
| np_dtype = getattr(np, dtype) |
| np.random.seed(0) |
| in_np = np.random.normal(0, 1.0 / C, (N, iH, C)).astype(np_dtype) |
| wt_np = np.random.normal(0, 1.0 / C, (O, kH, C)).astype(np_dtype) |
|
|
| in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| in_pt = torch.from_numpy(in_np.transpose(0, 2, 1)) |
| wt_pt = torch.from_numpy(wt_np.transpose(2, 0, 1)) |
|
|
| out_mx = mx.conv_transpose1d( |
| in_mx, |
| wt_mx, |
| stride=stride, |
| padding=padding, |
| output_padding=output_padding, |
| ) |
|
|
| out_pt = torch.conv_transpose1d( |
| in_pt, |
| wt_pt, |
| stride=stride, |
| padding=padding, |
| output_padding=output_padding, |
| ) |
| out_pt = torch.transpose(out_pt, 2, 1) |
|
|
| self.assertEqual(out_pt.shape, out_mx.shape) |
| self.assertTrue(np.allclose(out_pt.numpy(), out_mx, atol=atol)) |
|
|
| for dtype in ("float32",): |
| for N, C, O in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): |
| for iH, kH, stride, padding, output_padding in ( |
| (3, 2, 2, 0, 1), |
| (5, 3, 2, 1, 0), |
| (7, 4, 3, 1, 2), |
| ): |
| run_conv_transpose_1d_output_padding( |
| N, C, O, iH, kH, stride, padding, output_padding, dtype=dtype |
| ) |
|
|
| @unittest.skipIf(not has_torch, "requires Torch") |
| def test_torch_conv_transpose_2d_output_padding(self): |
| def run_conv_transpose_2d_output_padding( |
| N, |
| C, |
| O, |
| idim, |
| kdim, |
| stride, |
| padding, |
| output_padding, |
| dtype="float32", |
| atol=1e-5, |
| ): |
| with self.subTest( |
| dtype=dtype, |
| N=N, |
| C=C, |
| O=O, |
| idim=idim, |
| kdim=kdim, |
| stride=stride, |
| padding=padding, |
| output_padding=output_padding, |
| ): |
| np_dtype = getattr(np, dtype) |
| np.random.seed(0) |
| iH, iW = idim |
| kH, kW = kdim |
| in_np = np.random.normal(0, 1.0 / C, (N, iH, iW, C)).astype(np_dtype) |
| wt_np = np.random.normal(0, 1.0 / C, (O, kH, kW, C)).astype(np_dtype) |
|
|
| in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)) |
| wt_pt = torch.from_numpy(wt_np.transpose(3, 0, 1, 2)) |
|
|
| out_mx = mx.conv_transpose2d( |
| in_mx, |
| wt_mx, |
| stride=stride, |
| padding=padding, |
| output_padding=output_padding, |
| ) |
|
|
| out_pt = torch.conv_transpose2d( |
| in_pt, |
| wt_pt, |
| stride=stride, |
| padding=padding, |
| output_padding=output_padding, |
| ) |
| out_pt = torch.permute(out_pt, (0, 2, 3, 1)).numpy(force=True) |
|
|
| self.assertEqual(out_pt.shape, out_mx.shape) |
| self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) |
|
|
| for dtype in ("float32",): |
| for N, C, O in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): |
| for idim, kdim, stride, padding, output_padding in ( |
| ((3, 3), (2, 2), (2, 2), (0, 0), (1, 1)), |
| ((5, 5), (3, 3), (2, 2), (1, 1), (0, 0)), |
| ((7, 7), (4, 4), (3, 3), (1, 1), (2, 2)), |
| ): |
| run_conv_transpose_2d_output_padding( |
| N, |
| C, |
| O, |
| idim, |
| kdim, |
| stride, |
| padding, |
| output_padding, |
| dtype=dtype, |
| ) |
|
|
| @unittest.skipIf(not has_torch, "requires Torch") |
| def test_torch_conv_transpose_3d_output_padding(self): |
| def run_conv_transpose_3d_output_padding( |
| N, |
| C, |
| O, |
| idim, |
| kdim, |
| stride, |
| padding, |
| output_padding, |
| dtype="float32", |
| atol=1e-5, |
| ): |
| with self.subTest( |
| dtype=dtype, |
| N=N, |
| C=C, |
| O=O, |
| idim=idim, |
| kdim=kdim, |
| stride=stride, |
| padding=padding, |
| output_padding=output_padding, |
| ): |
| np_dtype = getattr(np, dtype) |
| np.random.seed(0) |
| iD, iH, iW = idim |
| kD, kH, kW = kdim |
| in_np = np.random.normal(0, 1.0 / C, (N, iD, iH, iW, C)).astype( |
| np_dtype |
| ) |
| wt_np = np.random.normal(0, 1.0 / C, (O, kD, kH, kW, C)).astype( |
| np_dtype |
| ) |
|
|
| in_mx, wt_mx = map(mx.array, (in_np, wt_np)) |
| in_pt = torch.from_numpy(in_np.transpose(0, 4, 1, 2, 3)) |
| wt_pt = torch.from_numpy(wt_np.transpose(4, 0, 1, 2, 3)) |
|
|
| out_mx = mx.conv_transpose3d( |
| in_mx, |
| wt_mx, |
| stride=stride, |
| padding=padding, |
| output_padding=output_padding, |
| ) |
| out_pt = torch.conv_transpose3d( |
| in_pt, |
| wt_pt, |
| stride=stride, |
| padding=padding, |
| output_padding=output_padding, |
| ) |
| out_pt = torch.permute(out_pt, (0, 2, 3, 4, 1)).numpy(force=True) |
|
|
| self.assertEqual(out_pt.shape, out_mx.shape) |
| self.assertTrue(np.allclose(out_pt, out_mx, atol=atol)) |
|
|
| for dtype in ("float32",): |
| for N, C, O in ((1, 1, 1), (1, 6, 1), (4, 32, 64)): |
| for idim, kdim, stride, padding, output_padding in ( |
| ((3, 3, 3), (2, 2, 2), (2, 2, 2), (0, 0, 0), (1, 1, 1)), |
| ((5, 5, 5), (3, 3, 3), (2, 2, 2), (1, 1, 1), (0, 0, 0)), |
| ((7, 7, 7), (4, 4, 4), (3, 3, 3), (1, 1, 1), (2, 2, 2)), |
| ): |
| run_conv_transpose_3d_output_padding( |
| N, |
| C, |
| O, |
| idim, |
| kdim, |
| stride, |
| padding, |
| output_padding, |
| dtype=dtype, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| mlx_tests.MLXTestRunner() |
|
|
