code: def test_correct_embed_out_with_diag_cov(self):
      n_harmonic_functions = 2
      x = torch.randn((1, 3))
      diag_cov = torch.randn((1, 3))
      D = 3 * n_harmonic_functions * 2  # sin + cos
  
      embed_fun = HarmonicEmbedding(
          n_harmonic_functions=n_harmonic_functions, append_input=False
      )
      embed_out = embed_fun(x, diag_cov=diag_cov)
  
      self.assertEqual(embed_out.shape, (1, D))
  
      # Compute the scaling factor introduce in MipNerf
      scale_factor = (
          -0.5
          * diag_cov[..., None]
          * torch.pow(embed_fun._frequencies[None, :], 2)
      )
      scale_factor = torch.exp(scale_factor).reshape(1, -1).tile((1, 2))
      # If we remove this scaling factor, we should go back to the
      # classical harmonic embedding:
      # Sum the squares of the respective frequencies
      # cos^2(x) + sin^2(x) = 1
      embed_out_without_cov = embed_out / scale_factor
      sum_squares = (
          embed_out_without_cov[0, : D // 2] ** 2
          + embed_out_without_cov[0, D // 2 :] ** 2
      )
      self.assertClose(sum_squares, torch.ones((D // 2)))
  
      # Test append input
      embed_fun = HarmonicEmbedding(
          n_harmonic_functions=n_harmonic_functions, append_input=True
      )
      embed_out_appended_input = embed_fun(x, diag_cov=diag_cov)
      self.assertClose(
          embed_out_appended_input.shape,
          torch.tensor((1, D + x.shape[-1])),
      )
      # Last plane in output is the input
      self.assertClose(embed_out_appended_input[..., -x.shape[-1] :], x)
      self.assertClose(
          embed_out_appended_input[..., : -x.shape[-1]], embed_out
      )
  
file_path: pytorch3d/tests/test_harmonic_embedding.py
imports: import unittest
  import torch
  from pytorch3d.renderer.implicit import HarmonicEmbedding
  from torch.distributions import MultivariateNormal
  from .common_testing import TestCaseMixin
label_desc: Use image node_id 5 for calling the TestHarmonicEmbedding obj's underlying member method code with example usage: obj.test_correct_embed_out_with_diag_cov() without return types