//  [python]
def test_weight_norm(self):
        def helper(dim, layer='linear', dtype=torch.float32):
            # linear layer
            if layer == 'linear':
                cpu_x = torch.randn((2, 5), device='cpu', dtype=dtype, requires_grad=True)
                x = cpu_x.detach().clone().to('mps').requires_grad_()

                cpu_weight = torch.randn(10, 5, device='cpu', dtype=dtype, requires_grad=True)
                weight = cpu_weight.detach().clone().to('mps').requires_grad_()

                cpu_bias = torch.randn(10, device='cpu', dtype=dtype, requires_grad=True)
                bias = cpu_bias.detach().clone().to('mps').requires_grad_()

                cpu_linear = torch.nn.Linear(5, 10, device='cpu')
                linear = torch.nn.Linear(5, 10, device='mps')

                with torch.no_grad():
                    cpu_linear.weight.copy_(cpu_weight)
                    cpu_linear.bias.copy_(cpu_bias)
                    linear.weight.copy_(weight)
                    linear.bias.copy_(bias)

                cpu_norm = torch.nn.utils.weight_norm(cpu_linear, dim=dim)
                norm = torch.nn.utils.weight_norm(linear, dim=dim)

                cpu_out = cpu_norm(cpu_x)
                out = norm(x)

                self.assertEqual(cpu_out, out)

                cpu_grad = torch.randn(cpu_out.shape)
                grad = cpu_grad.to('mps')
                cpu_out.backward(gradient=cpu_grad)
                out.backward(gradient=grad)

                self.assertEqual(cpu_linear.weight_g.grad, linear.weight_g.grad)
                self.assertEqual(cpu_linear.weight_v.grad, linear.weight_v.grad)

                self.assertEqual(x.grad, cpu_x.grad)

            # conv layer
            if layer == 'conv':
                cpu_x = torch.randn((3, 5, 5), device='cpu', dtype=dtype, requires_grad=True)
                x = cpu_x.detach().clone().to('mps').requires_grad_()

                cpu_conv = torch.nn.Conv2d(3, 3, 3, device='cpu')
                conv = torch.nn.Conv2d(3, 3, 3, device='mps')

                with torch.no_grad():
                    conv.weight.copy_(cpu_conv.weight)
                    conv.bias.copy_(cpu_conv.bias)

                cpu_norm = torch.nn.utils.weight_norm(cpu_conv, dim=dim)
                norm = torch.nn.utils.weight_norm(conv, dim=dim)

                cpu_out = cpu_norm(cpu_x)
                out = norm(x)

                self.assertEqual(cpu_out, out)

                cpu_grad = torch.randn(cpu_out.shape)
                grad = cpu_grad.to('mps')
                cpu_out.backward(gradient=cpu_grad)
                out.backward(gradient=grad)

                self.assertEqual(cpu_conv.weight_g.grad, conv.weight_g.grad)
                self.assertEqual(cpu_conv.weight_v.grad, conv.weight_v.grad)

                self.assertEqual(x.grad, cpu_x.grad)

                # conv layer
            if layer == 'conv3d':
                cpu_x = torch.randn((3, 5, 5, 4), device='cpu', dtype=dtype, requires_grad=True)
                x = cpu_x.detach().clone().to('mps').requires_grad_()

                cpu_conv = torch.nn.Conv3d(3, 3, 3, device='cpu')
                conv = torch.nn.Conv3d(3, 3, 3, device='mps')

                with torch.no_grad():
                    conv.weight.copy_(cpu_conv.weight)
                    conv.bias.copy_(cpu_conv.bias)

                cpu_norm = torch.nn.utils.weight_norm(cpu_conv, dim=dim)
                norm = torch.nn.utils.weight_norm(conv, dim=dim)

                cpu_out = cpu_norm(cpu_x)
                out = norm(x)

                self.assertEqual(cpu_out, out)

                cpu_grad = torch.randn(cpu_out.shape)
                grad = cpu_grad.to('mps')
                cpu_out.backward(gradient=cpu_grad)
                out.backward(gradient=grad)

                self.assertEqual(cpu_conv.weight_g.grad, conv.weight_g.grad)
                self.assertEqual(cpu_conv.weight_v.grad, conv.weight_v.grad)

                self.assertEqual(x.grad, cpu_x.grad)

        helper(0, layer='linear')
        helper(1, layer='linear')
        helper(-1, layer='linear')

        helper(0, layer='conv')
        helper(1, layer='conv')
        helper(2, layer='conv')
        helper(3, layer='conv')
        helper(-1, layer='conv')

        if product_version >= 13.2:
            # Conv3d is only available from MacOS 13 onwards
            helper(0, layer='conv3d')
            helper(1, layer='conv3d')
            helper(2, layer='conv3d')
            helper(3, layer='conv3d')
            helper(4, layer='conv3d')
            helper(-1, layer='conv3d')