torchrec icon indicating copy to clipboard operation
torchrec copied to clipboard

KJT dynamo path for conditions

Open IvanKobzarev opened this issue 1 year ago • 4 comments

Summary: The change is no-change for eager path.

Dynamo path

Dynamo tracing can not pass through various conditions => making dynamo path without those conditions.

We skip asserts for dynamo and variable batch check if it can go through "equal batches" logic => for Dynamo fallback on Variable Batch path.

Comments about torch scripting

guard_size_oblivious(), _check_is_size() is not torch scriptable atm, we have to guard it in is not torch.jit.is_scripting

torch.jit._unwrap_optional(recat) is used to hint torch script that recat is not None at the places where we use it.

torch.jit.annotate() hints to TorchScript about the returned type of tensor.tolist()

Differential Revision: D55695198

IvanKobzarev avatar Apr 03 '24 14:04 IvanKobzarev

This pull request was exported from Phabricator. Differential Revision: D55695198

facebook-github-bot avatar Apr 03 '24 14:04 facebook-github-bot

This pull request was exported from Phabricator. Differential Revision: D55695198

facebook-github-bot avatar Apr 03 '24 14:04 facebook-github-bot

This pull request was exported from Phabricator. Differential Revision: D55695198

facebook-github-bot avatar Apr 03 '24 14:04 facebook-github-bot

This pull request was exported from Phabricator. Differential Revision: D55695198

facebook-github-bot avatar Apr 03 '24 14:04 facebook-github-bot