refactor autoencoder tests (temporal decoder, cosmos, kvae, mochi)

[akshan-main]

What does this PR do?

Part of the ongoing modeling-test migration (following #13369 and #13153). Migrates the AutoencoderKLTemporalDecoder, AutoencoderKLCosmos, AutoencoderKLKVAE, and AutoencoderKLMochi test suites to the mixin-based structure (Config + ModelTesterMixin / TrainingTesterMixin / MemoryTesterMixin / NewAutoencoderTesterMixin).

Also fixes AutoencoderKLMochi.forward to propagate return_dict to decode, which was returning a DecoderOutput instead of a tensor with return_dict=False.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? (Discussed on Slack with @DN6.)
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@sayakpaul@DN6

[sayakpaul]

@askserge could you do a review?

[github-actions]

🤗 Serge says:

Overall

Clean migration of four autoencoder test suites to the new mixin-based structure. The pattern is consistent with other already-migrated tests in the repo. The model-side fix to AutoencoderKLMochi.forward is correct and well-motivated.

Correctness

• Mochi forward fix — The old code called self.decode(z) (always return_dict=True), then manually checked if not return_dict: return (dec,), which would wrap a DecoderOutput in a tuple instead of returning a raw tensor tuple. The new code correctly propagates return_dict to decode, which already handles both branches. Good fix.

• Dropped test_sharded_checkpoints_device_map skip — The old Mochi test skipped this test (with a multi-GPU error reason). The new ModelTesterMixin (from tests/models/testing_utils/common.py) does not define test_sharded_checkpoints_device_map, so the skip is correctly dropped.

• Renamed test_effective_gradient_checkpointing → test_gradient_checkpointing_equivalence — Matches the method name in the new TrainingTesterMixin. Correct.

• input_shape property removed — Not used anywhere in the new mixin framework or NewAutoencoderTesterMixin. Correct to drop.

Minor observations

• The generator property creates a new torch.Generator on every access. This is intentional (ensures fresh state per call) and matches the pattern in other migrated tests (test_models_autoencoder_kl.py, test_models_autoencoder_dc.py, etc.).

• The main_input_name property is only consumed by MemoryTesterMixin (in test_layerwise_casting_training). All four configs correctly define it.

• Pre-existing: the return type annotation on AutoencoderKLMochi.forward is torch.Tensor | torch.Tensor (line 1095, not in this diff) — should probably be DecoderOutput | tuple[torch.Tensor], but that's out of scope for this PR.

20 LLM turns · 28 tool calls · 89.2s · 439893 in / 4044 out tokens

[akshan-main]

@sayakpaul
every other VAE uses

DecoderOutput | torch.Tensor

return type annotation

I'll make that change too

[sayakpaul]

Thanks!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.