code: def rfl2norm2(xf, xs, axis=(0, 1)):
      r
  
      scl = 1.0 / np.prod(np.array([xs[k] for k in axis]))
      slc0 = (slice(None),) * axis[-1]
      nrm0 = np.linalg.norm(xf[slc0 + (0,)])
      idx1 = (xs[axis[-1]] + 1) // 2
      nrm1 = np.linalg.norm(xf[slc0 + (slice(1, idx1),)])
      if xs[axis[-1]] % 2 == 0:
          nrm2 = np.linalg.norm(xf[slc0 + (slice(-1, None),)])
      else:
          nrm2 = 0.0
      return scl*(nrm0**2 + 2.0*nrm1**2 + nrm2**2)
docs: r"""
      Compute the squared :math:`\ell_2` norm in the DFT domain, taking
      into account the unnormalised DFT scaling, i.e. given the DFT of a
      multi-dimensional array computed via :func:`rfftn`, return the
      squared :math:`\ell_2` norm of the original array.
  
      Parameters
      ----------
      xf : array_like
        Input array
      xs : sequence of ints
        Shape of original array to which :func:`rfftn` was applied to
        obtain the input array
      axis : sequence of ints, optional (default (0,1))
        Axes on which the input is in the frequency domain
  
      Returns
      -------
      x : float
        :math:`\|\mathbf{x}\|_2^2` where the input array is the result of
        applying :func:`rfftn` to the specified axes of multi-dimensional
        array :math:`\mathbf{x}`