Add native masked MSE loss for Sapiens2ForPoseEstimation
Author: Sainava
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #46518

As discussed with the maintainers in the linked issue, this PR implements native supervised pose-estimation loss directly in `Sapiens2ForPoseEstimation` to unlock fine-tuning capabilities using the `Trainer` API.

**Specific Changes:**
* Added an optional `target_weights` parameter to the `forward` signature to handle keypoint visibility masking, following the masking behavior of OpenMMLab's `KeypointMSELoss` without adding external dependencies.
* Implemented the masked MSE loss calculation using pure PyTorch, explicitly casting masks to the heatmap dtype to ensure safe `fp16`/`bf16` mixed-precision training.
* Added strict shape and dimensionality validation for supervision targets to prevent silent broadcasting errors.
* Updated the `forward` docstring to explicitly document the new parameter.
* Updated `test_modeling_sapiens2.py` to verify the loss computation both with and without `target_weights`.

All local `make fix-repo` and `make check-repo` checks have passed.

**Scope of this PR**

This implementation follows the masking behavior of OpenMMLab's `KeypointMSELoss` through optional `target_weights` support. It intentionally does not implement `skip_empty_channel` or the configurable `loss_weight` parameter from OpenMMLab, as neither is currently exposed through the Sapiens2 Transformers configuration. The goal of this PR is to provide native supervised pose-estimation training support while keeping the initial implementation focused and aligned with existing Transformers conventions.

<!--
## Code Agent Policy

The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. We are currently bottlenecked by our ability to review and respond to them. As a result, 
**we ask that new users do not submit pure code agent PRs** at this time. 
You may use code agents in drafting or to help you diagnose issues. We'd also ask autonomous "OpenClaw"-like agents
not to open any PRs or issues for the moment.

PRs that appear to be fully agent-written will probably be closed without review, and we may block users who do this
repeatedly or maliciously. 

This is a rapidly-evolving situation that's causing significant shockwaves in the open-source community. As a result, 
this policy is likely to be updated regularly in the near future. For more information, please read [`CONTRIBUTING.md`](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md).
-->
- [X] I confirm that this is not a pure code agent PR.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://huggingface.co/docs/transformers/contributing) and the
      [Pull Request](https://huggingface.co/docs/transformers/pr_checks) checks?
- [X] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case. https://github.com/huggingface/transformers/issues/46518#issuecomment-4743822498
- [X] Did you make sure to update the documentation with your changes according to the [guidelines](https://github.com/huggingface/transformers/tree/main/docs)?
- [X] Did you write any new necessary [tests](https://huggingface.co/docs/transformers/testing)?


## Who can review?

Hi @guarin, following up on the discussion in #46518, I've put together a draft implementation for supervised pose-estimation loss support in Sapiens2ForPoseEstimation. I'd appreciate any feedback when you have a chance.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez @vasqu
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @Cyrilvallez (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc
- kernels: @vasqu @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Comments:
Hi @guarin, thanks for the feedback!

I've updated the PR to follow the suggested loss-function architecture, moved the pose-estimation test inputs into a dedicated factory method, and added a supervised fine-tuning example to the documentation.

One implementation detail: `_loss_function` is wrapped with `staticmethod(...)` so it works correctly with the `self.loss_function` property and avoids Python binding the model instance as the first argument.

 I'd really appreciate any further feedback when you have a chance :)


Hi @guarin, thanks for the review! I've pushed the suggested updates (tests and documentation)

Regarding the loss function: I tried passing `weight=label_weights` directly to `self.loss_function`, but since the default implementation resolves to `torch.nn.functional.mse_loss`, this raises:

```text
ValueError: Weights and input must have the same size
```

when using OpenMMLab-style visibility weights of shape `[batch_size, num_keypoints, 1, 1]` against heatmaps of shape `[batch_size, num_keypoints, height, width]`.

To preserve the original masking behavior, I'm currently computing the unreduced loss and applying the broadcasted weights explicitly.

Would you prefer that I instead move the weighting logic into a small custom `_loss_function` implementation so the forward pass can remain:

```python
loss = self.loss_function(heatmaps, labels, weight=label_weights)
```

while preserving the same behavior?

And yes, I'd definitely be interested in working on the preprocessing side as well. My preference would be to get the training-loss support merged first and then tackle the preprocessing changes in a follow-up PR if that's okay .

Hi! I think we can safely assume that `label_weights` has the same shape as `labels` in the `forward` function. The pre-processing can make sure that this is the case and expand any `(batch_size, num_keypoints, 1, 1)` tensors to match `labels.shape` so we don't have to worry about that in `forward`. 

**CI Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=46764)

Hi @guarin ! I've simplified the forward pass to use weight=label_weights directly in the loss function as suggested.
Thanks again for the guidance :)

Hi @vasqu, thanks for the review :) .I've moved the loss to LOSS_MAPPING and added the .backward() tests as well .

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: sapiens2

run-slow: sapiens2

[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/28160283990)

This comment contains `run-slow`, running the specified jobs: 

models: ["models/sapiens2"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/28160283990)

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [8bce9d31](https://github.com/huggingface/transformers/commit/8bce9d31451ad48f7c4d10a712700b3e37e3c2d3) | workflow commit (merge commit) |
| PR | [ffb1c626](https://github.com/Sainava/Sai-transformers/commit/ffb1c62699821a0194fe235d2e56e045a4abf818) | branch commit (from PR) |
| main | [cf284d4f](https://github.com/huggingface/transformers/commit/cf284d4f6d72ffb095ce4c8fda51d19714c8f9ec) | base commit (on `main`) |

✅ No failing test specific to this PR 🎉 👏 !


The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46764). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.