Make LazyAwaitable wait_async dynamo traceable

pytorch / torchrec

Pytorch domain library for recommendation systems

BSD 3-Clause "New" or "Revised" License

1.79k stars 381 forks source link

Make LazyAwaitable wait_async dynamo traceable #2184

Open IvanKobzarev opened 1 week ago

IvanKobzarev commented 1 week ago

Summary: Clients requested to be able to call directly LazyAwaitable._wait_async(a) under compiled region.

LazyAwaitable.wait_async uses torch.fx.node.map_aggregate which is skipped by dynamo => copied inside torchrec, changed immutable_list -> list, immutable_dict -> dict (immutable versions were failing in dynamo).

Using only for compiled path, not to break any BC on eager path.

Differential Revision: D59113305

facebook-github-bot commented 1 week ago

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