From 6622b18386b7ac4196dccc70537f7ab03abc74fc Mon Sep 17 00:00:00 2001
From: Vijayraj Gohil <63923064+vraj130@users.noreply.github.com>
Date: Thu, 16 Jul 2026 11:42:22 -0700
Subject: [PATCH] Fix crash in TensorFlow GNM compute_vertex_normals.

gnm_tensorflow.GNM.compute_vertex_normals crashes on any real mesh
with:

  InvalidArgumentError: data.shape = [35324,3] does not start with
  segment_ids.shape = [1,105972] [Op:UnsortedSegmentSum]

The face normals tensor (B, T, 3) was reshaped straight to (B*T, 3)
without replicating each face normal for the three corner vertices it
contributes to, while the segment ids enumerate all B*T*3 triangle
corners. tf.math.unsorted_segment_sum requires segment_ids.shape to
be a prefix of data.shape, so the shapes can never match for a real
mesh. The numpy, jax and pytorch backends all replicate the face
normal across the three corners (e.g. np.add.at broadcasting in
gnm_numpy.py); the tensorflow backend was the only one missing this
step, and the only one without test coverage for normals.

Fix: tile the face normals across the corner axis before flattening,
and flatten the batch-offset segment ids to match, reproducing the
numpy backend's semantics.

Add test_compute_vertex_normals_numpy_parity to gnm_tensorflow_test,
checking TensorFlow normals against the numpy implementation (itself
validated against tensorflow_graphics in gnm_numpy_test) over
unbatched and multi-dimensional batch shapes. The test reproduces the
InvalidArgumentError before this fix and passes after; the full
gnm_tensorflow_test module passes.
---
 gnm/shape/gnm_tensorflow.py      |  8 ++++++--
 gnm/shape/gnm_tensorflow_test.py | 26 ++++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/gnm/shape/gnm_tensorflow.py b/gnm/shape/gnm_tensorflow.py
index ec7833f2..0a55c11b 100644
--- a/gnm/shape/gnm_tensorflow.py
+++ b/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
     )
diff --git a/gnm/shape/gnm_tensorflow_test.py b/gnm/shape/gnm_tensorflow_test.py
index 2244d926..17934eb9 100644
--- a/gnm/shape/gnm_tensorflow_test.py
+++ b/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,),