//  [python]
def test_norm_preserved_relu(self):

        # Check if the norm of the features is approximately preserved

        # To do so, we try to apply ReLU over strictly positive features
        # or apply ReLU over roughly centered features (expecting the norm to reduce by approx sqrt(2))

        # for these tests to work, we should not apply normalization of the Fourier matrices, otherwise the scale
        # of the output features depends on the number of Fourier components recovered

        g = no_base_space(so3_group(1))

        for F, N in zip(range(1, 4), [24, 120, 300]):
            print(F, N)

            cl = FourierPointwise(g, 1, g.fibergroup.bl_irreps(F), normalize=False, function='p_relu', type='thomson', N=N)

            c = cl.in_type.size
            B = 256

            # make sure the average value of the feature is very high so the ReLU acts as an identity
            # then, check that the norm of the input and output features are similar
            x = torch.randn(B, c)
            x = x.view(B, len(cl.in_type), cl.rho.size)
            p = 0
            for irr in cl.rho.irreps:
                irr = g.irrep(*irr)
                if irr.is_trivial():
                    x[:, :, p] = 150.
                p += irr.size

            in_norms = torch.linalg.norm(x, dim=2).reshape(-1).cpu().detach().numpy()
            x = x.view(B, cl.in_type.size)
            x = cl.in_type(x).transform_fibers(g.fibergroup.sample())
            y = cl(x).tensor
            y = y.view(B, len(cl.out_type), cl.rho_out.size)
            out_norms = torch.linalg.norm(y, dim=2).reshape(-1).cpu().detach().numpy()
            self.assertTrue(np.allclose(in_norms, out_norms))

            # now, make sure features are centered around zero such that, on average, half of the entries are set to 0
            # by ReLU. We expect the output to have a norm roughly sqrt(2) smaller than the input
            x = torch.randn(B, c)
            x = x.view(B, len(cl.in_type), cl.rho.size)
            p = 0
            for irr in cl.rho.irreps:
                irr = g.irrep(*irr)
                if irr.is_trivial():
                    x[:, :, p] = 0.
                p += irr.size

            in_norms = torch.linalg.norm(x, dim=2).reshape(-1).cpu().detach().numpy()
            x = x.view(B, cl.in_type.size)
            x = cl.in_type(x).transform_fibers(g.fibergroup.sample())
            y = cl(x).tensor
            y = y.view(B, len(cl.out_type), cl.rho_out.size)
            out_norms = torch.linalg.norm(y, dim=2).reshape(-1).cpu().detach().numpy()
            norms_ratio = out_norms / in_norms
            print(norms_ratio.max(), norms_ratio.min(), norms_ratio.mean(), norms_ratio.std())
            self.assertTrue(np.allclose(norms_ratio, 1./np.sqrt(2), atol=1e-1, rtol=1e-1), msg=f"{np.fabs(norms_ratio-1./np.sqrt(2)).max()}")

        for F, N in zip(range(1, 4), [16, 34, 66]):
            print(F, N)

            grid = g.fibergroup.sphere_grid(type='thomson', N=N)
            cl = QuotientFourierPointwise(g, (False, -1), 40, g.fibergroup.bl_irreps(F), normalize=False, function='p_relu', grid=grid)

            c = cl.in_type.size
            B = 32

            # make sure the average value of the feature is very high so the ReLU acts as an identity
            # then, check that the norm of the input and output features are similar
            x = torch.randn(B, c)
            x = x.view(B, len(cl.in_type), cl.rho.size)
            p = 0
            for irr in cl.rho.irreps:
                irr = g.irrep(*irr)
                if irr.is_trivial():
                    x[:, :, p] = 150.
                p += irr.size
            x = x.view(B, cl.in_type.size)

            in_norms = torch.linalg.norm(x, dim=1).reshape(-1).cpu().detach().numpy()
            x = cl.in_type(x).transform_fibers(g.fibergroup.sample())
            y = cl(x).tensor
            out_norms = torch.linalg.norm(y, dim=1).reshape(-1).cpu().detach().numpy()
            self.assertTrue(np.allclose(in_norms, out_norms))

            # now, make sure features are centered around zero such that, on average, half of the entries are set to 0
            # by ReLU. We expect the output to have a norm roughly sqrt(2) smaller than the input
            x = torch.randn(B, c)
            x = x.view(B, len(cl.in_type), cl.rho.size)
            p = 0
            for irr in cl.rho.irreps:
                irr = g.irrep(*irr)
                if irr.is_trivial():
                    x[:, :, p] = 0.
                p += irr.size

            x = x.view(B, cl.in_type.size)
            in_norms = torch.linalg.norm(x, dim=1).reshape(-1).cpu().detach().numpy()
            x = cl.in_type(x).transform_fibers(g.fibergroup.sample())
            y = cl(x).tensor
            out_norms = torch.linalg.norm(y, dim=1).reshape(-1).cpu().detach().numpy()
            norms_ratio = out_norms / in_norms
            # print(norms_ratio.max(), norms_ratio.min(), norms_ratio.mean(), norms_ratio.std())
            self.assertTrue(np.allclose(norms_ratio, 1. / np.sqrt(2), atol=1e-1, rtol=1e-1), msg=f"{np.fabs(norms_ratio-1./np.sqrt(2)).max()}")

        max_f = 5
        for F, N in zip(range(1, max_f+1), [200]*max_f):
            print(F, N)

            grid = g.fibergroup.sphere_grid(type='thomson', N=N)
            cl = QuotientFourierPointwise(g, (False, -1), 30, g.fibergroup.bl_irreps(1), out_irreps=g.fibergroup.bl_irreps(F), normalize=False, function='p_relu', grid=grid)

            c = cl.in_type.size
            B = 32

            # make sure the average value of the feature is very high so the ReLU acts as an identity
            # then, check that the norm of the input and output features are similar
            x = torch.randn(B, c)
            x = x.view(B, len(cl.in_type), cl.rho.size)
            p = 0
            for irr in cl.rho.irreps:
                irr = g.irrep(*irr)
                if irr.is_trivial():
                    x[:, :, p] = 150.
                p += irr.size
            x = x.view(B, cl.in_type.size)

            in_norms = torch.linalg.norm(x, dim=1).reshape(-1).cpu().detach().numpy()
            x = cl.in_type(x).transform_fibers(g.fibergroup.sample())
            y = cl(x).tensor
            out_norms = torch.linalg.norm(y, dim=1).reshape(-1).cpu().detach().numpy()
            self.assertTrue(np.allclose(in_norms, out_norms))

            # now, make sure features are centered around zero such that, on average, half of the entries are set to 0
            # by ReLU. We expect the output to have a norm roughly sqrt(2) smaller than the input
            x = torch.randn(B, c)
            x = x.view(B, len(cl.in_type), cl.rho.size)
            p = 0
            for irr in cl.rho.irreps:
                irr = g.irrep(*irr)
                if irr.is_trivial():
                    x[:, :, p] = 0.
                p += irr.size

            x = x.view(B, cl.in_type.size)
            in_norms = torch.linalg.norm(x, dim=1).reshape(-1).cpu().detach().numpy()
            x = cl.in_type(x).transform_fibers(g.fibergroup.sample())
            y = cl(x).tensor
            out_norms = torch.linalg.norm(y, dim=1).reshape(-1).cpu().detach().numpy()
            norms_ratio = out_norms / in_norms
            self.assertTrue(np.allclose(norms_ratio, 1. / np.sqrt(2), atol=1e-1, rtol=1e-1), msg=f"{np.fabs(norms_ratio-1./np.sqrt(2)).max()}")