Repository: huggingface/transformers
Pull request: fix(models): Apply STE in Dac.from_latents to match the forward pass

### What does this PR do?

The following issue was identified and fixed in this PR:

→ Implemented a fix that applies the straight-through estimator to each latent chunk before `out_proj`, following the straight-through estimator pattern used in the forward pass ([`DacResidualVectorQuantizer.from_latents`](https://github.com/huggingface/transformers/blob/main/src/transformers/models/dac/modeling_dac.py#L371-L400) was missing  STE before calling `out_proj`). Also fixes some CI regressions, still not sure which commit caused them, but CI is a tad bit happier again :)
→ Added the repro script to the linked issue.

Fixes #43819.

**Before the fix (feel free to cross-check; these errors are reproducible):**

<img width="500" height="600" alt="before_dac_2" src="https://github.com/user-attachments/assets/5234dd2e-3117-41d0-a09d-3d308f39930b" />

<img width="600" height="500" alt="before_dac" src="https://github.com/user-attachments/assets/23bf6f4c-e4b4-4171-9a00-a13307460080" />

**After the fix (feel free to cross-check):**

<img width="500" height="600" alt="after_dac_2" src="https://github.com/user-attachments/assets/03e94a7d-c36d-4e91-aa0c-0b3d304024a5" />

<img width="600" height="500" alt="after_dac" src="https://github.com/user-attachments/assets/3063ac25-767d-449c-a797-3586881c964e" />

### 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://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [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.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you fix any necessary existing tests?

Changed files:
filename: src/transformers/models/dac/modeling_dac.py
patch: @@ -390,11 +390,13 @@ def from_latents(self, latents: torch.Tensor):
         n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[0]
         for i in range(n_codebooks):
             hidden_dim_j, hidden_dim_k = dims[i], dims[i + 1]
-            quantized_latents_i, codes_i = self.quantizers[i].decode_latents(latents[:, hidden_dim_j:hidden_dim_k, :])
+            latent_chunk = latents[:, hidden_dim_j:hidden_dim_k, :]
+            quantized_latents_i, codes_i = self.quantizers[i].decode_latents(latent_chunk)
             quantized_latents.append(quantized_latents_i)
             codes.append(codes_i)
 
-            quantized_representation_i = self.quantizers[i].out_proj(quantized_latents_i)
+            quantized_with_ste = latent_chunk + (quantized_latents_i - latent_chunk)
+            quantized_representation_i = self.quantizers[i].out_proj(quantized_with_ste)
             quantized_representation = quantized_representation + quantized_representation_i
 
         return quantized_representation, torch.cat(quantized_latents, dim=1)

filename: tests/models/dac/test_modeling_dac.py
patch: @@ -802,6 +802,7 @@ def test_quantizer_from_latents_integration(self, model_name):
         with torch.no_grad():
             encoder_outputs = model.encode(input_values)
             latents = encoder_outputs.projected_latents
+            original_quantizer_representation = encoder_outputs.quantized_representation
 
             # reconstruction using from_latents
             quantizer_representation, quantized_latents = model.quantizer.from_latents(latents=latents)
@@ -810,6 +811,11 @@ def test_quantizer_from_latents_integration(self, model_name):
             # forward pass
             original_reconstructed = model(input_values).audio_values
 
+        # ensure quantizer representations match
+        self.assertTrue(
+            torch.allclose(quantizer_representation, original_quantizer_representation, atol=1e-6),
+            msg="Quantizer representation from from_latents should match original quantizer forward pass",
+        )
         # ensure forward and decode are the same
         self.assertTrue(
             torch.allclose(reconstructed, original_reconstructed, atol=1e-6),