Fix equivolume batch broadcasting

kaolin/metrics/tetmesh.py

@@ -61,15 +61,15 @@ def equivolume(tet_vertices, tetrahedrons_mean=None, pow=4):
         tetrahedrons_mean (torch.Tensor):
             Mean volume of all tetrahedrons in a grid,
             of shape :math:`(\text{batch_size})` or :math:`(1,)` (broadcasting).
-            Default: Compute ``torch.mean(tet_vertices, dim=-1)``.
+            Default: Compute the mean tetrahedron volume for each mesh.
         pow (int):
             Power for the equivolume loss.
             Increasing power puts more emphasis on the larger tetrahedron deformation.
             Default: 4.
 
     Returns:
         (torch.Tensor):
-            EquiVolume loss for each mesh, of shape :math:`(\text{batch_size})`.
+            EquiVolume loss for each mesh, of shape :math:`(\text{batch_size}, 1)`.
 
     Example:
         >>> tet_vertices = torch.tensor([[[[0.5000, 0.5000, 0.7500],
@@ -89,8 +89,8 @@ def equivolume(tet_vertices, tetrahedrons_mean=None, pow=4):
         ...                                [0.6000, 0.9000, 0.3000],
         ...                                [0.5500, 0.3500, 0.9000]]]])
         >>> equivolume(tet_vertices, pow=4)
-        tensor([[2.2961e-10],
-                [7.7704e-10]])
+        tensor([[2.2898e-15],
+                [1.5422e-09]])
     """
     _validate_tet_vertices(tet_vertices)
 
@@ -100,7 +100,7 @@ def equivolume(tet_vertices, tetrahedrons_mean=None, pow=4):
     if tetrahedrons_mean is None:
         # finding the mean volume of all tetrahedrons in the tetrahedron grid
         tetrahedrons_mean = torch.mean(volumes, dim=-1)
-    tetrahedrons_mean = tetrahedrons_mean.reshape(1, -1)
+    tetrahedrons_mean = tetrahedrons_mean.reshape(-1, 1)
     # compute EquiVolume loss
     equivolume_loss = torch.mean(torch.pow(
         torch.abs(volumes - tetrahedrons_mean), exponent=pow),

tests/python/kaolin/metrics/test_tetmesh.py

@@ -76,4 +76,15 @@ def test_equivolume(self):
                                        [0.9000, 0.8000, 0.7000],
                                        [0.6000, 0.9000, 0.3000],
                                        [0.5500, 0.3500, 0.9000]]]])
-        assert torch.allclose(tetmesh.equivolume(tetrahedrons, pow=4), torch.tensor([[2.2898e-15], [2.9661e-10]]))
+        expected = torch.tensor([[2.2898e-15], [1.5422e-09]])
+        assert torch.allclose(tetmesh.equivolume(tetrahedrons, pow=4), expected, rtol=1e-4, atol=1e-18)
+
+    def test_equivolume_batched_meshes(self):
+        heights = torch.tensor([[6., 18., 30.], [60., 84., 108.]])
+        tetrahedrons = torch.zeros((2, 3, 4, 3))
+        tetrahedrons[..., 1, 0] = 1.
+        tetrahedrons[..., 2, 1] = 1.
+        tetrahedrons[..., 3, 2] = heights
+
+        expected = torch.tensor([[32. / 3.], [512. / 3.]])
+        assert torch.allclose(tetmesh.equivolume(tetrahedrons, pow=4), expected)