🤗 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

Thanks!