Skip to content

Fix crash in TensorFlow GNM compute_vertex_normals.#28

Open
vraj130 wants to merge 1 commit into
google:mainfrom
vraj130:fix-tf-vertex-normals
Open

Fix crash in TensorFlow GNM compute_vertex_normals.#28
vraj130 wants to merge 1 commit into
google:mainfrom
vraj130:fix-tf-vertex-normals

Conversation

@vraj130

@vraj130 vraj130 commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

gnm_tensorflow.GNM.compute_vertex_normals crashes on any real mesh:

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

(Reproduced on the bundled V3 head model: the forward pass succeeds, then
computing normals for the resulting (1, 17821, 3) vertices raises the above.)

Root cause: the face-normals tensor (B, T, 3) is 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 (T > 1).
The numpy, jax, and pytorch backends all replicate the face normal across the
three corners (e.g. the np.add.at broadcast in gnm_numpy.py); the
TensorFlow backend was the only one missing this step and the only backend
with no test coverage for compute_vertex_normals.

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

Test: adds test_compute_vertex_normals_numpy_parity, mirroring the
existing test_parity_with_gnm_numpy pattern: TensorFlow normals must match
the numpy implementation (itself validated against tensorflow_graphics in
gnm_numpy_test.py) over unbatched and multi-dimensional batch shapes
((), (1,), (1,2), (2,1,2)). The new test fails with the
InvalidArgumentError above before this fix and passes after; the full
gnm_tensorflow_test module passes.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant