Add integrated position encoding based on MIPNerf implementation.

Summary: Add a new implicit module Integral Position Encoding based on [MIP-NeRF](https://arxiv.org/abs/2103.13415).

Reviewed By: shapovalov

Differential Revision: D46352730

fbshipit-source-id: c6a56134c975d80052b3a11f5e92fd7d95cbff1e

tests/implicitron/test_implicit_function_neural_radiance_field.py

@@ -0,0 +1,66 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from pytorch3d.implicitron.models.implicit_function.base import ImplicitronRayBundle
+from pytorch3d.implicitron.models.implicit_function.neural_radiance_field import (
+    NeuralRadianceFieldImplicitFunction,
+)
+
+
+class TestNeuralRadianceFieldImplicitFunction(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(42)
+
+    def test_forward_with_integrated_positionial_embedding(self):
+        shape = [2, 4, 4]
+        ray_bundle = ImplicitronRayBundle(
+            origins=torch.randn(*shape, 3),
+            directions=torch.randn(*shape, 3),
+            bins=torch.randn(*shape, 6 + 1),
+            lengths=torch.randn(*shape, 6),
+            pixel_radii_2d=torch.randn(*shape, 1),
+            xys=None,
+        )
+        model = NeuralRadianceFieldImplicitFunction(
+            n_hidden_neurons_dir=32, use_integrated_positional_encoding=True
+        )
+        raw_densities, ray_colors, _ = model(ray_bundle=ray_bundle)
+
+        self.assertEqual(raw_densities.shape, (*shape, ray_bundle.lengths.shape[-1], 1))
+        self.assertEqual(ray_colors.shape, (*shape, ray_bundle.lengths.shape[-1], 3))
+
+    def test_forward_with_integrated_positionial_embedding_raise_exception(self):
+        shape = [2, 4, 4]
+        ray_bundle = ImplicitronRayBundle(
+            origins=torch.randn(*shape, 3),
+            directions=torch.randn(*shape, 3),
+            bins=None,
+            lengths=torch.randn(*shape, 6),
+            pixel_radii_2d=torch.randn(*shape, 1),
+            xys=None,
+        )
+        model = NeuralRadianceFieldImplicitFunction(
+            n_hidden_neurons_dir=32, use_integrated_positional_encoding=True
+        )
+        with self.assertRaises(ValueError):
+            _ = model(ray_bundle=ray_bundle)
+
+    def test_forward(self):
+        shape = [2, 4, 4]
+        ray_bundle = ImplicitronRayBundle(
+            origins=torch.randn(*shape, 3),
+            directions=torch.randn(*shape, 3),
+            lengths=torch.randn(*shape, 6),
+            pixel_radii_2d=torch.randn(*shape, 1),
+            xys=None,
+        )
+        model = NeuralRadianceFieldImplicitFunction(n_hidden_neurons_dir=32)
+        raw_densities, ray_colors, _ = model(ray_bundle=ray_bundle)
+        self.assertEqual(raw_densities.shape, (*shape, 6, 1))
+        self.assertEqual(ray_colors.shape, (*shape, 6, 3))