Fix and update type hints for `make_functional.py` #91579

XuehaiPan · 2023-01-02T12:35:32Z

Changes in details:

Fix and update some out-of-date type hints in _functorch/make_functional.py.
~~Explicitly use OrderedDict for order-sensitive mappings.~~

In create_names_map(), _swap_state(), and FunctionalModuleWithBuffers.__init__(), the unordered dict was used. The key order should be preserved for dict.items() while it is required to zip with a tuple of params/buffers. Although since Python 3.6, the built-in dictionary is insertion ordered (PEP 468). Explicit is better than implicit.

cc @zou3519

pytorch-bot · 2023-01-02T12:35:34Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/91579

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit bce866f:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

linux-foundation-easycla · 2023-01-02T12:35:36Z

The committers listed above are authorized under a signed CLA.

✅ login: XuehaiPan / name: Xuehai Pan (524f59d)

torch/_functorch/make_functional.py

zou3519

Thanks for the PR, @XuehaiPan. We should try to enable mypy checking for this file in PyTorch CI (to check the type hints and prevent regressions -- thanks for catching incorrect type annotations).

I don't buy using OrderedDict over dict here -- if you feel strongly about it, could we please leave it to a separate discussion?

zou3519 · 2023-01-03T16:29:24Z

torch/_functorch/make_functional.py

+    named_params = OrderedDict(named_params)
+    tied_named_params = OrderedDict(tied_named_params)


I don't buy the dict -> OrderedDict change:

Performance. OrderedDict has a different implementation than dict and there isn't a need to construct a new OrderedDict around a dict here.

As you mentioned, dict() is already ordered. As of Python 3.7, the language standard says that dict() is ordered. PyTorch does not support Python < 3.7, so there are no BC concerns.

consistency: I'd like to see buy in for "prefer OrderedDict to dict" in other parts of the PyTorch codebase.

I don't buy the dict -> OrderedDict change:

Performance. OrderedDict has a different implementation than dict and there isn't a need to construct a new OrderedDict around a dict here.

As you mentioned, dict() is already ordered. As of Python 3.7, the language standard says that dict() is ordered. PyTorch does not support Python < 3.7, so there are no BC concerns.

Let's pass on this since the language standard guarantees the insertion order for dict since Python 3.7. The builtin dict is faster for iteration.

In [1]: from collections import OrderedDict In [2]: iterable = tuple((str(i), i) for i in range(100)) In [3]: %timeit dict(iterable) 2.37 µs ± 20.5 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each) In [4]: %timeit OrderedDict(iterable) 8.42 µs ± 86.5 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each) In [5]: d = dict(iterable) In [6]: od = OrderedDict(iterable) In [7]: %timeit list(d.items()) 2.41 µs ± 83.5 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each) In [8]: %timeit list(od.items()) 4.05 µs ± 87.4 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

consistency: I'd like to see buy in for "prefer OrderedDict to dict" in other parts of the PyTorch codebase.

pytorch/torch/nn/modules/module.py

Lines 448 to 463 in 3120054

super().__setattr__('training', True)

super().__setattr__('_parameters', OrderedDict())

super().__setattr__('_buffers', OrderedDict())

super().__setattr__('_non_persistent_buffers_set', set())

super().__setattr__('_backward_pre_hooks', OrderedDict())

super().__setattr__('_backward_hooks', OrderedDict())

super().__setattr__('_is_full_backward_hook', None)

super().__setattr__('_forward_hooks', OrderedDict())

super().__setattr__('_forward_hooks_with_kwargs', OrderedDict())

super().__setattr__('_forward_pre_hooks', OrderedDict())

super().__setattr__('_forward_pre_hooks_with_kwargs', OrderedDict())

super().__setattr__('_state_dict_hooks', OrderedDict())

super().__setattr__('_state_dict_pre_hooks', OrderedDict())

super().__setattr__('_load_state_dict_pre_hooks', OrderedDict())

super().__setattr__('_load_state_dict_post_hooks', OrderedDict())

super().__setattr__('_modules', OrderedDict())

The attributes in nn.Module, such as _paramters, _buffers, _modules are stored as OrderedDict.

Note that, del and set attribute in _swap_state changes the order of items in _paramters and _buffers. (del + set is equivalent to move_to_end())

pytorch/torch/_functorch/make_functional.py

Lines 134 to 142 in 3120054

def _swap_state(mod: nn.Module, names_map: List[str], elems):

result = []

for (_, attr_names), elem in zip(names_map.items(), elems):

for i, attr_name in enumerate(attr_names):

if i == 0:

result.append(_get_nested_attr(mod, attr_name))

_del_nested_attr(mod, attr_name)

_set_nested_attr(mod, attr_name, elem)

return result

zou3519

thank you!

XuehaiPan · 2023-01-16T12:45:39Z

@zou3519 Hi, any update for this PR? I opened a follow-up PR in #92218 which may share some relevant changes.

kshitij12345 · 2023-01-18T07:49:34Z

Once approved by a PyTorch maintainer, you can trigger the merge yourself by commenting @pytorchbot merge.

More ref about the pytorchbot: https://github.com/pytorch/pytorch/wiki/Bot-commands

XuehaiPan · 2023-01-18T11:51:53Z

@pytorchbot merge

pytorchmergebot · 2023-01-18T11:53:36Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2023-01-18T11:53:37Z

Merge failed

Reason: This PR is too stale; the last push date was more than 3 days ago. Please rebase and try again. You can rebase by leaving the following comment on this PR:
@pytorchbot rebase

Details for Dev Infra team

Raised by workflow job

XuehaiPan · 2023-01-18T11:53:51Z

@pytorchbot rebase

XuehaiPan · 2023-01-18T11:55:09Z

@pytorchbot merge

pytorchmergebot · 2023-01-18T11:55:41Z

@pytorchbot successfully started a rebase job. Check the current status here

pytorchmergebot · 2023-01-18T11:55:46Z

Successfully rebased make-functional onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout make-functional && git pull --rebase)

pytorchmergebot · 2023-01-18T11:56:51Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2023-01-18T17:55:29Z

The merge job was canceled. If you believe this is a mistake,then you can re trigger it through pytorch-bot.

XuehaiPan · 2023-01-18T18:23:53Z

@pytorchbot merge

pytorchmergebot · 2023-01-18T18:25:49Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

…#92536) - Fixes #92295 - Resolves #86708 - Resolves #92153 - Closes #92401 - Closes #92218 - Requires #91579 Refactor NN stateless APIs by swapping module tensors. Pull Request resolved: #92536 Approved by: https://github.com/jbschlosser

pytorchbot added the open source label Jan 2, 2023

XuehaiPan force-pushed the make-functional branch 2 times, most recently from 2afefe2 to 2329d7a Compare January 2, 2023 13:15

zou3519 reviewed Jan 3, 2023

View reviewed changes

torch/_functorch/make_functional.py Show resolved Hide resolved

zou3519 reviewed Jan 3, 2023

View reviewed changes

XuehaiPan force-pushed the make-functional branch 5 times, most recently from 9ea0f0e to dd122d2 Compare January 4, 2023 10:47

zou3519 approved these changes Jan 4, 2023

View reviewed changes

XuehaiPan mentioned this pull request Jan 15, 2023

[functorch] Cache intermediate submodules for submodule lookup in stateless functional calls #92218

Closed

XuehaiPan mentioned this pull request Jan 18, 2023

[functorch][nn] Refactor NN stateless APIs by swapping module tensors #92536

Closed

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jan 18, 2023

XuehaiPan added 4 commits January 18, 2023 11:55

Fix and update type hints for make_functional.py

93fdf51

Prefer dict over OrderedDict

3c2a701

Enable linting for make_functional

e6b880b

Format make_functional.py

bce866f

pytorchmergebot force-pushed the make-functional branch from dd122d2 to bce866f Compare January 18, 2023 11:55

pytorchmergebot added the Merged label Jan 18, 2023

pytorchmergebot closed this in ae4ec7d Jan 18, 2023

XuehaiPan deleted the make-functional branch January 18, 2023 19:25

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix and update type hints for `make_functional.py` #91579

Fix and update type hints for `make_functional.py` #91579

XuehaiPan commented Jan 2, 2023 •

edited

pytorch-bot bot commented Jan 2, 2023 •

edited

linux-foundation-easycla bot commented Jan 2, 2023 •

edited

zou3519 left a comment

zou3519 Jan 3, 2023

XuehaiPan Jan 4, 2023

zou3519 left a comment

XuehaiPan commented Jan 16, 2023

kshitij12345 commented Jan 18, 2023

XuehaiPan commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

XuehaiPan commented Jan 18, 2023

XuehaiPan commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

XuehaiPan commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

		named_params = OrderedDict(named_params)
		tied_named_params = OrderedDict(tied_named_params)

	super().__setattr__('training', True)
	super().__setattr__('_parameters', OrderedDict())
	super().__setattr__('_buffers', OrderedDict())
	super().__setattr__('_non_persistent_buffers_set', set())
	super().__setattr__('_backward_pre_hooks', OrderedDict())
	super().__setattr__('_backward_hooks', OrderedDict())
	super().__setattr__('_is_full_backward_hook', None)
	super().__setattr__('_forward_hooks', OrderedDict())
	super().__setattr__('_forward_hooks_with_kwargs', OrderedDict())
	super().__setattr__('_forward_pre_hooks', OrderedDict())
	super().__setattr__('_forward_pre_hooks_with_kwargs', OrderedDict())
	super().__setattr__('_state_dict_hooks', OrderedDict())
	super().__setattr__('_state_dict_pre_hooks', OrderedDict())
	super().__setattr__('_load_state_dict_pre_hooks', OrderedDict())
	super().__setattr__('_load_state_dict_post_hooks', OrderedDict())
	super().__setattr__('_modules', OrderedDict())

	def _swap_state(mod: nn.Module, names_map: List[str], elems):
	result = []
	for (_, attr_names), elem in zip(names_map.items(), elems):
	for i, attr_name in enumerate(attr_names):
	if i == 0:
	result.append(_get_nested_attr(mod, attr_name))
	_del_nested_attr(mod, attr_name)
	_set_nested_attr(mod, attr_name, elem)
	return result

Fix and update type hints for make_functional.py #91579

Fix and update type hints for make_functional.py #91579

Conversation

XuehaiPan commented Jan 2, 2023 • edited

pytorch-bot bot commented Jan 2, 2023 • edited

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/91579

✅ No Failures

linux-foundation-easycla bot commented Jan 2, 2023 • edited

zou3519 left a comment

Choose a reason for hiding this comment

zou3519 Jan 3, 2023

Choose a reason for hiding this comment

XuehaiPan Jan 4, 2023

Choose a reason for hiding this comment

zou3519 left a comment

Choose a reason for hiding this comment

XuehaiPan commented Jan 16, 2023

kshitij12345 commented Jan 18, 2023

XuehaiPan commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

Merge started

pytorchmergebot commented Jan 18, 2023

Merge failed

XuehaiPan commented Jan 18, 2023

XuehaiPan commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

Merge started

pytorchmergebot commented Jan 18, 2023

XuehaiPan commented Jan 18, 2023

pytorchmergebot commented Jan 18, 2023

Merge started

Fix and update type hints for `make_functional.py` #91579

Fix and update type hints for `make_functional.py` #91579

XuehaiPan commented Jan 2, 2023 •

edited

pytorch-bot bot commented Jan 2, 2023 •

edited

linux-foundation-easycla bot commented Jan 2, 2023 •

edited