TAPNext coordinate cross entropy ignores visibility weights after scalar reduction

Open
  • Is something wrong?
#168 0 comments 0 reactions 0 assignees View on GitHub

@sylvesterkaczmarek is already working on this.

Since Oct 9, 2026.

  • #169 by @sylvesterkaczmarek — closed without merging
  • #170 by @sylvesterkaczmarek — open

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
35/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Stale
Tech stack
python, pytorch
Domain
computer-vision, machine-learning, testing

Research direction

Start by reading tapnet/tapnext/torch_losses.py, focusing on coordinate_softmax and how its cross-entropy reductions interact with the visibility mask. Run the prepared regressions comparing against the independent logsumexp/gather calculation and exercising tapnext_loss_and_grad; done means masked points contribute neither coordinate loss nor gradients while both mask shapes remain supported. A linked pull request is already open.

Written by the indexing model from the issue text.

Description

Reproduced on main 730cda1c730877cfedbe01bf87fb1cadb78a565d.

coordinate_softmax uses the default mean reduction in both F.cross_entropy calls before multiplying by the visibility mask. The resulting scalar is broadcast over the mask, so the mask cancels out whenever its sum is nonzero. Occluded points still affect coordinate loss and gradients.

import torch
from tapnet.tapnext import torch_losses
logits = torch.tensor([[[[4., 0., 4., 0.], [0., 4., 0., 4.]]]], requires_grad=True)
labels = torch.full((1, 1, 2, 2), .5)
mask = torch.tensor([[[[1.], [0.]]]])
loss = torch_losses.coordinate_softmax(logits, labels, mask, pixel_size=2)
loss.backward()
print(loss.item(), logits.grad)

Current loss is approximately 4.0363 rather than the visible-point value 0.0363, and the masked point receives nonzero gradients. The public tapnext_loss_and_grad path is also affected.

Use unreduced cross entropy and then apply per-point weights. Preserve support for both plain masks and the trailing singleton axis used by the training helper. Reference for unreduced output: https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html

Prepared regressions compare an independent logsumexp/gather calculation and exercise the actual public backward helper with a deterministic tiny tracker. All-zero masks and nonfinite logits are outside this fix.

Dominant language
Jupyter Notebook
Stars
2.2k
Forks
215
No merged PRs in 30d

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from google-deepmind/tapnet

All issues in google-deepmind/tapnet

Similar issues

More Computer Vision issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.