Fix crash in TensorFlow GNM compute_vertex_normals.

gnm/shape/gnm_tensorflow.py

@@ -125,14 +125,18 @@ def compute_vertex_normals(
     v2 = face_vertices[..., 2, :]
     face_normals_area = tf.experimental.numpy.cross(v1 - v0, v2 - v0, axis=-1)
 
-    data = tf.reshape(face_normals_area[:, :, None, :], (-1, 3))
+    # Replicate each face normal for the three corner vertices it contributes
+    # to, matching the flattened (triangle, corner) order of the segment ids.
+    data = tf.reshape(
+        tf.tile(face_normals_area[:, :, None, :], [1, 1, 3, 1]), (-1, 3)
+    )
     segment_ids = tf.reshape(self.triangles, (-1,))
     segment_ids = tf.tile(segment_ids[None, :], [batch_size, 1])
 
     batch_offset = (
         tf.range(batch_size, dtype=segment_ids.dtype)[:, None] * num_vertices
     )
-    flat_segment_ids = segment_ids + batch_offset
+    flat_segment_ids = tf.reshape(segment_ids + batch_offset, (-1,))
     flat_normals = tf.math.unsorted_segment_sum(
         data, flat_segment_ids, num_segments=batch_size * num_vertices
     )

gnm/shape/gnm_tensorflow_test.py

@@ -108,6 +108,32 @@ def test_parity_with_gnm_numpy(
 
     np.testing.assert_almost_equal(actual, desired, decimal=4)
 
+  @parameterized.product(
+      version=_MAINTAINED_MAJOR_GNM_VERSIONS,
+      variant=tuple(_SUPPORTED_VARIANTS),
+      batch_dims=BATCH_DIMS,
+  )
+  def test_compute_vertex_normals_numpy_parity(
+      self, version: str, variant: str, batch_dims: tuple[int, ...]
+  ):
+    """Tests that TensorFlow vertex normals match the NumPy implementation."""
+    if variant not in self.gnms_np[version]:
+      self.skipTest(f'variant {variant} not supported in {version}.')
+    gnm_np = self.gnms_np[version][variant]
+    gnm_tf = self.gnms_tf[version][variant]
+
+    parameters_np = gnm_test_utils.random_gnm_parameters(
+        gnm_np, batch_shape=batch_dims, seed=self.rng
+    )
+    vertices_np = gnm_np(**parameters_np)
+    desired = gnm_np.compute_vertex_normals(vertices_np)
+
+    vertices_tf = tf.convert_to_tensor(vertices_np, dtype=tf.float32)
+    actual = gnm_tf.compute_vertex_normals(vertices_tf)
+
+    self.assertEqual(tuple(actual.shape), vertices_np.shape)
+    np.testing.assert_allclose(actual.numpy(), desired, atol=1e-4)
+
   @parameterized.product(
       version=_MAINTAINED_MAJOR_GNM_VERSIONS,
       variant=(gnm_tensorflow.GNMVariant.HEAD,),