//  [python]
def test_groupNorm_mixed_dtype(self):
        def helper(size, groups, memory_format, dtype):
            channels = size[1]
            input = torch.randn(size).cpu().to(dtype=dtype)
            input_bf1 = (
                input.contiguous(memory_format=memory_format)
                .detach()
                .requires_grad_(True)
            )
            input_bf2 = input_bf1.clone().detach().requires_grad_(True)
            input_f = input_bf1.float().detach().requires_grad_(True)
            m_bf = nn.GroupNorm(groups, channels).cpu().to(dtype=dtype)
            m_f = deepcopy(m_bf).float()
            m_f2 = deepcopy(m_f)
            # bfloat16 input and bfloat16 parameters
            out = m_bf(input_bf1)
            # bfloat16 input and float parameters
            out2 = m_f(input_bf2)
            # float input and float parameters
            out3 = m_f2(input_f)
            torch.testing.assert_close(out, out2, atol=5e-3, rtol=5e-3)
            torch.testing.assert_close(out2.float(), out3, atol=5e-3, rtol=5e-3)
            grad_out = torch.randn(out2.shape).cpu().to(dtype=dtype)
            grad_out_bf1 = (
                grad_out.contiguous(memory_format=memory_format)
                .detach()
                .requires_grad_(True)
            )
            grad_out_bf2 = grad_out_bf1.clone().detach().requires_grad_(True)
            grad_out_f = grad_out_bf2.clone().float().detach().requires_grad_(True)
            # bfloat16 input grad and float parameters
            out2.backward(grad_out_bf2, retain_graph=True)
            # float input grad and float parameters
            out3.backward(grad_out_f, retain_graph=True)
            # bfloat16 input grad and bfloat16 parameters
            out.backward(grad_out_bf1, retain_graph=True)
            torch.testing.assert_close(
                m_f.weight.grad, m_f2.weight.grad, atol=2e-5, rtol=1e-5
            )
            torch.testing.assert_close(
                m_f.bias.grad, m_f2.bias.grad, atol=1e-5, rtol=1e-5
            )
            torch.testing.assert_close(
                input_bf2.grad.float(), input_f.grad, atol=5e-5, rtol=5e-3
            )
            # full bf16 has lower precision compared with mixed bf16 and fp32.
            atol = None
            rtol = None
            if dtype == torch.bfloat16:
                atol = 1e-2
                rtol = 1.3e-1
            else:
                assert dtype == torch.half
                atol = 5e-3
                rtol = 1.5e-2
            torch.testing.assert_close(
                m_bf.weight.grad.float(), m_f.weight.grad, atol=atol, rtol=rtol
            )
            torch.testing.assert_close(
                m_bf.bias.grad.float(), m_f.bias.grad, atol=atol, rtol=rtol
            )
            torch.testing.assert_close(
                input_bf1.grad, input_bf2.grad, atol=atol, rtol=rtol
            )

        cl_formats = {4: torch.channels_last, 5: torch.channels_last_3d}
        for dtype in [torch.bfloat16, torch.half]:
            for shape, g in [
                ((1, 8, 4, 3), 2),
                ((1, 8, 3, 4), 4),
                ((4, 40, 40, 40), 2),
                ((4, 8, 40, 40), 4),
                ((1, 8, 40, 40), 4),
                ((1, 8, 40, 40), 2),
                ((1, 8, 50, 50), 2),
                ((1, 8, 50, 50), 4),
                ((1, 40, 50, 50), 2),
                ((1, 9, 3, 4, 5), 3),
                ((1, 60, 10, 10, 10), 3),
                ((1, 9, 10, 50, 50), 3),
                ((1, 60, 10, 50, 50), 3),
                ((1, 8, 65, 55), 2),
                ((1, 3, 65, 55), 1),
                ((1, 3, 20, 20), 1),
            ]:
                for is_cl in [False, True]:
                    format = (
                        cl_formats[len(shape)] if is_cl else torch.contiguous_format
                    )
                    helper(shape, g, format, dtype)