Skip to content

Align TBE fake batch shape with its operator contract (#4481) - #4481

Open
kaanbaloglu wants to merge 1 commit into
meta-pytorch:mainfrom
kaanbaloglu:export-D114760345
Open

Align TBE fake batch shape with its operator contract (#4481)#4481
kaanbaloglu wants to merge 1 commit into
meta-pytorch:mainfrom
kaanbaloglu:export-D114760345

Conversation

@kaanbaloglu

@kaanbaloglu kaanbaloglu commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

Summary:

ir_tbe_lookup_impl returns output tensors using the supplied batch_size, but its fake registration discarded that value and invented an unbacked dimension. This change makes the fake use the supplied batch_size, matching the real implementation and the sibling ir_emb_lookup and ir_kt_regroup fakes.

This is a shape-contract consistency cleanup, not a claim of new runtime batch polymorphism. For static exports, it makes the lookup interior batch dimension input-derived instead of unbacked. Dynamic exports may prove the same relation through downstream constraints. ir_dynamic_batch_emb_lookup remains the dedicated operator for cases that intentionally require an unbacked batch dimension.

The BUCK target-label updates are AUTODEPS2-required normalization for the two touched Python targets; no dependency behavior change is intended.

Differential Revision: D114760345

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 5, 2026
@meta-codesync

meta-codesync Bot commented Aug 5, 2026

Copy link
Copy Markdown
Contributor

@kaanbaloglu has exported this pull request. If you are a Meta employee, you can view the originating Diff in D114760345.

)

Summary:

`ir_tbe_lookup_impl` returns output tensors using the supplied `batch_size`, but its fake registration discarded that value and invented an unbacked dimension. This change makes the fake use the supplied `batch_size`, matching the real implementation and the sibling `ir_emb_lookup` and `ir_kt_regroup` fakes.

This is a shape-contract consistency cleanup, not a claim of new runtime batch polymorphism. For static exports, it makes the lookup interior batch dimension input-derived instead of unbacked. Dynamic exports may prove the same relation through downstream constraints. `ir_dynamic_batch_emb_lookup` remains the dedicated operator for cases that intentionally require an unbacked batch dimension.

The BUCK target-label updates are AUTODEPS2-required normalization for the two touched Python targets; no dependency behavior change is intended.

Differential Revision: D114760345
@meta-codesync meta-codesync Bot changed the title Align TBE fake batch shape with its operator contract Align TBE fake batch shape with its operator contract (#4481) Aug 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant