Fix imports for fake op wrappers used in export

[geoffreyQiu]

Description

• Add imports for fake op wrappers used for torch.export
• Add directives to avoid formatter removing/reordering the import

[greptile-apps]

Greptile Summary

This PR fixes missing imports for fake op wrapper modules (dynamicemb meta modules, hstu_cuda_ops, fake_hstu_cuda_ops, and hstu.hstu_ops_gpu) required for torch.export to work correctly. It uses # isort: off/on guards in exportable_embedding.py to preserve the load-order dependency between the real op registration and the fake impl registration.

Confidence Score: 5/5

Safe to merge — changes are additive import fixes with no logic modifications.

Both files receive only import additions. The load-order invariant (real ops before fake impls) is correctly preserved via isort guards in exportable_embedding.py. In fused_hstu_op.py the new hstu.hstu_ops_gpu import is a submodule of hstu, so Python's package initialization guarantees the parent's init.py runs first regardless of line order, making an isort guard unnecessary there. No logic changes, no regressions.

No files require special attention.

Important Files Changed

Filename	Overview
examples/hstu/modules/exportable_embedding.py	Adds isort-protected import block registering fake op implementations (dynamicemb meta modules, hstu_cuda_ops, fake_hstu_cuda_ops) needed for torch.export; also repositions the section divider comment.
examples/hstu/ops/fused_hstu_op.py	Adds import of hstu.hstu_ops_gpu immediately after import hstu to register fake op implementations for torch.export; no isort guard added but import order is naturally safe due to Python submodule initialization semantics.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Module import starts] --> B[_load_inference_emb_ops\nloads inference_emb_ops.so\nregisters torch.ops.INFERENCE_EMB.*]
    B --> C[import dynamicemb.index_range_meta\nregisters fake impls for torch.export]
    C --> D[import dynamicemb.lookup_meta\nregisters fake impls for torch.export]
    D --> E[import hstu_cuda_ops\nregisters torch.ops.hstu_cuda_ops.*]
    E --> F[import commons.ops.cuda_ops.fake_hstu_cuda_ops\nregisters fake impls for torch.export]
    F --> G[isort: on - Normal imports resume]
    G --> H[ExportableEmbedding class available]

    subgraph fused_hstu_op.py
        I[import hstu\nregisters torch.ops.fbgemm.*] --> J[import hstu.hstu_ops_gpu\nregisters fake impls for torch.export]
    end

Loading

_{Reviews (1): Last reviewed commit: "Fix imports for fake ops wrapper used in..." | Re-trigger Greptile}

[JacoCheung]

/build

[JacoCheung]

/build

[JacoCheung]

❌ Pipeline #48606612 -- failed

Job	Status	Log
pre_check	✅ success	view
train_build	✅ success	view
inference_build	✅ success	view
tritonserver_build	✅ success	view
build_whl	✅ success	view
dynamicemb_test_fwd_bwd_8gpus	✅ success	view
dynamicemb_test_load_dump_8gpus	✅ success	view
unit_test_1gpu_a100	❌ failed	view
unit_test_1gpu_h100	❌ failed	view
unit_test_4gpu	❌ failed	view
unit_test_tp_4gpu	❌ failed	view
L20_unit_test_1gpu	✅ success	view
inference_unit_test_1gpu	❌ failed	view
inference_test_1gpu	❌ failed	view

Result: 8/14 jobs passed

View full pipeline

[JacoCheung]

Close as meged to #363