Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[export] use tree_map for _flatten_dynamic_shapes #125415

Closed
wants to merge 1 commit into from
Closed

[export] use tree_map for _flatten_dynamic_shapes #125415

wants to merge 1 commit into from

Commits on May 2, 2024

  1. [export] use tree_map for _flatten_dynamic_shapes (#125415) 

    Summary:
Pull Request resolved: #125415

Fixing the implementation of `_flatten_dynamic_shapes()`, to follow how `_process_dynamic_shapes()` does it. The previous implementation would misinterpret some nested dynamic shapes specs, causing it to miss out on some shapes specs, for example with nested inputs/constant input tuples:

```
inputs = (
    (2, 1),
    (
        torch.randn(2, 1),
        torch.randn(2, 2),
        torch.randn(2, 3),
    )
)

dynamic_shapes = (
    (None, None),
    (
        None,
        None,
        None,
    )
)
```
This would get interpreted as 2 shapes specs for 2d and 3d tensors. Fixing so this doesn't happen.

Test Plan: Existing export tests

Differential Revision: D56894923
    @pianpwk @facebook-github-bot
    pianpwk authored and facebook-github-bot committed May 2, 2024
    Configuration menu
    Copy the full SHA
    5c9ddd1 View commit details
    Browse the repository at this point in the history