[FEATURE] TensorDictModule in_keys allowed as Dict[str, tuple | list] to enable multi use of a sample feature #1101

bachdj-px · 2024-11-21T17:11:49Z

Description

See feature request #1099

Motivation and Context

Close issue #1099

I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

New feature (non-breaking change which adds core functionality)
Documentation (update in the documentation)
Example inside the code documentation

Checklist

I have read the CONTRIBUTION guide (required)
My change requires a change to the documentation.
I have updated the tests accordingly (required for a bug fix or a new feature).
I have updated the documentation accordingly.

facebook-github-bot · 2024-11-21T17:11:54Z

Hi @bachdj-px!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!

vmoens

Hello, thanks for this PR!

IIUC your change is to allow one element of the input tensordict to be used multiple times within the module inputs.

That would be a good addition.

I think the tests would need to be a bit more extensive. We also need to check this works with other TD modules, like ProbabilisticTensorDictModule.

tensordict/nn/common.py

vmoens · 2024-11-23T12:24:34Z

test/test_nn.py

+    def test_tuple_input_keys(self):
+        in_keys = {"1": "first", "2": ("second", "third")}
+        fn = lambda x, y, z: x + y + z
+        module = TensorDictModule(fn, in_keys=in_keys, out_keys=["out"])
+
+        kword_in_keys = sorted(zip(module._kwargs, module.in_keys))
+        assert kword_in_keys == [("first", "1"), ("second", "2"), ("third", "2")]


I don't really get this example:
the modue is not run, what would the input tensordict look like?
Could the example be a bit more extensive?

Thanks, it's a good point - I updated that test. I also added another test that tests the integration of TensorDictModule within a ProbabilisticTensorDictSequential.

@vmoens Thanks for your feedback.
Could you clarify how you would see the change with ProbabilisticTensorDictModule? It seems to me that this feature will only impact TensorDictModule. ProbabilisticTensorDictModule inherits from TensorDictModuleBase and is not affected. Also I don't see the use of implementing my suggested feature in ProbabilisticTensorDictModule - its input have to be unique as far I see it.
Please let me know your thoughts

Sure
What if I want to do something along the line of

ProbabilisticTensorDictModule( in_keys={"hidden": ("concentration", "rate")}, out_keys=["other"], distribution_class=torch.distributions.Gamma )

with the gamma distribution taking params concentration and rate.

As per the utility of doing such thing - I think it's up to the users. To me as developer is more a matter of keeping the API as uniform as we can. You could for instance decide to build 2 different Normal distributions with a CompositeDistribution class that share the same scale?

Thanks for the explanations @vmoens .
Sorry, I am still a bit confused - the API documentation states that If in_keys is a dictionary, the keys are the keys of the distribution and the values are the keys in the tensordict that will get match to the corresponding distribution keys.. In your example, it looks to me that you should pass the in_keys dictionary in the reverse order. Passing the same tensordict key to 2 different distribution parameters seems to be already supported, for instance via:

ProbabilisticTensorDictModule( in_keys={"concentration": "value_in_input_td", "rate": "value_in_input_td" }, out_keys=["other"], distribution_class=torch.distributions.Gamma )

We could change the API to match your example but I suppose that it would be a breaking API change, wouldn't it?

Yeah i just realized that, so we are asking dicts that are flipped between TensorDictModule and ProbabilisticTensorDictModule?
That ain't good!
If we flip the dicts in TensorDictModule (to match Proba), maybe that would solve your problem altogether?

Absolutely, that would change the problem altogether! And it's a relatively change to make (docs and tests I assume)
However, that would be a breaking change for users already using dictionary with TensorDictModule. How would you like to proceed?

Yep that's BC breaking but it's for the best
Here is how I would procede in this case: what we want is that people passing in_keys now with a dict don't see things breaking but they get a warning, and we also want to make your life easy if you're using the new API.
So I'd suggest to add a out_to_in_map: bool | None = None keyword arg in the module that says:

if None, the current default is used but if in_keys is a dict, then a deprec warning is raised

If False we use the current behaviour input-key: output-key, no warning is raised.

If True, we use the new behaviour.

The warning must be explicit and tell the user that the default will change 2 versions from now (next is 0.7 so it will be in 0.9). To silence the warning pass out_to_in_map=False or True.

Does that makes sense? Should I take care of it?

Yes, thanks, I like this approach, it's safe. I can push the change today and add some tests for the 3 cases (check the warning is raised etc)

tensordict/nn/common.py

vmoens

Thanks! Looks great! I'd just make the doc / warning a bit more informative.

I don't think we should mention in the docstrings that the value can be None as this is just used internally to check if the user passed a value or not (documenting it as public API defies its purpose IMO).

vmoens · 2024-11-27T16:05:57Z

tensordict/nn/common.py

+            if out_to_in_map is None:
+                warnings.warn(
+                    "Using a dictionary in_keys without specifying out_to_in_map is deprecated."
+                    "Use out_to_in_map to indicate the ordering of the input keys.",


I think we want to change the behavior - currently out_to_in_map is False by default but we would like to make it True eventually. This should be mentioned in the warning, as well as the version where the change will occur (v0.9).

vmoens · 2024-11-27T16:07:10Z

tensordict/nn/common.py

        out_keys (iterable of str): keys to be written to the input tensordict. The length of out_keys must match the
            number of tensors returned by the embedded module. Using "_" as a key avoid writing tensor to output.

    Keyword Args:
+        out_to_in_map (bool or None, optional): if ``True``, `in_keys` is read as if the keys are the arguments keys of


Suggested change

out_to_in_map (bool or None, optional): if ``True``, `in_keys` is read as if the keys are the arguments keys of

out_to_in_map (bool, optional): if ``True``, `in_keys` is read as if the keys are the arguments keys of

vmoens · 2024-11-27T16:09:42Z

tensordict/nn/common.py

+            `False`, keys are considered go be the input keys and values the method's arguments. If `None` (default), the
+            behaviour is the same as for `False` but a deprecation warning is raised.


Suggested change

`False`, keys are considered go be the input keys and values the method's arguments. If `None` (default), the

behaviour is the same as for `False` but a deprecation warning is raised.

`False`, keys are considered go be the input keys and values the method's arguments. Deaults to ``False``.

.. warning:: The default value of `out_to_in_map` will change from `False` to `True` in the v0.9 release.

bachdj-px · 2024-11-27T19:08:45Z

Thanks for your feedback @vmoens ! I applied your suggestions - please let me know your thoughts
I am running the tests locally but many CUDA-related fail - they arise from test_tensordict.py and they seem not to be linked to the new feature. Could we run the CI to double check?

vmoens

Almost there thanks so much

vmoens · 2024-11-28T15:40:26Z

tensordict/nn/common.py

@@ -773,11 +773,19 @@ class TensorDictModule(TensorDictModuleBase):
            order given by the in_keys iterable.
            If ``in_keys`` is a dictionary, its keys must correspond to the key
            to be read in the tensordict and its values must match the name of
-            the keyword argument in the function signature.
+            the keyword argument in the function signature. If `out_to_in_map` is True,


Suggested change

the keyword argument in the function signature. If `out_to_in_map` is True,

the keyword argument in the function signature. If `out_to_in_map` is ``True``,

vmoens · 2024-11-28T15:40:59Z

tensordict/nn/common.py

+                The default value of `out_to_in_map` will change from `False` to `True` in the v0.9 release.
        inplace (bool or string, optional): if ``True`` (default), the output of the module are written in the tensordict


Suggested change

The default value of `out_to_in_map` will change from `False` to `True` in the v0.9 release.

inplace (bool or string, optional): if ``True`` (default), the output of the module are written in the tensordict

The default value of `out_to_in_map` will change from `False` to `True` in the v0.9 release.

inplace (bool or string, optional): if ``True`` (default), the output of the module are written in the tensordict

I also added the double quotes around False and True

tensordict/nn/common.py

… flexible input key dispatching

vmoens · 2024-11-29T14:49:47Z

Thanks for this!

facebook-github-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 Nov 22, 2024

vmoens reviewed Nov 23, 2024

View reviewed changes

bachdj-px force-pushed the support_tuple_in_keys_dispatch branch from 0af7100 to 198404e Compare November 25, 2024 13:38

vmoens added the enhancement New feature or request label Nov 26, 2024

bachdj-px force-pushed the support_tuple_in_keys_dispatch branch 3 times, most recently from 9e188f5 to 25b3f07 Compare November 26, 2024 19:15

vmoens reviewed Nov 27, 2024

View reviewed changes

bachdj-px force-pushed the support_tuple_in_keys_dispatch branch from 25b3f07 to 49a0201 Compare November 27, 2024 18:07

bachdj-px force-pushed the support_tuple_in_keys_dispatch branch from 49a0201 to c9e5be5 Compare November 28, 2024 15:18

vmoens reviewed Nov 28, 2024

View reviewed changes

feat(nn.common.TensorDictModule): support tuple values in in_keys for…

a35825b

… flexible input key dispatching

bachdj-px force-pushed the support_tuple_in_keys_dispatch branch from c9e5be5 to a35825b Compare November 28, 2024 15:54

vmoens merged commit e871b7d into pytorch:main Nov 29, 2024
47 of 52 checks passed

vmoens mentioned this pull request Dec 2, 2024

[Feature Request] Support tuple values in TensorDictModule in_keys arguments. #1099

Closed

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[FEATURE] TensorDictModule in_keys allowed as Dict[str, tuple | list] to enable multi use of a sample feature #1101

[FEATURE] TensorDictModule in_keys allowed as Dict[str, tuple | list] to enable multi use of a sample feature #1101

bachdj-px commented Nov 21, 2024

facebook-github-bot commented Nov 21, 2024

vmoens left a comment

vmoens Nov 23, 2024

bachdj-px Nov 25, 2024

bachdj-px Nov 25, 2024

vmoens Nov 26, 2024

bachdj-px Nov 26, 2024

vmoens Nov 26, 2024

bachdj-px Nov 26, 2024

vmoens Nov 26, 2024

bachdj-px Nov 26, 2024

vmoens left a comment

vmoens Nov 27, 2024

vmoens Nov 27, 2024

vmoens Nov 27, 2024

bachdj-px commented Nov 27, 2024

vmoens left a comment

vmoens Nov 28, 2024

vmoens Nov 28, 2024

bachdj-px Nov 28, 2024

vmoens commented Nov 29, 2024

	out_to_in_map (bool or None, optional): if ``True``, `in_keys` is read as if the keys are the arguments keys of
	out_to_in_map (bool, optional): if ``True``, `in_keys` is read as if the keys are the arguments keys of

		`False`, keys are considered go be the input keys and values the method's arguments. If `None` (default), the
		behaviour is the same as for `False` but a deprecation warning is raised.

	the keyword argument in the function signature. If `out_to_in_map` is True,
	the keyword argument in the function signature. If `out_to_in_map` is ``True``,

		The default value of `out_to_in_map` will change from `False` to `True` in the v0.9 release.
		inplace (bool or string, optional): if ``True`` (default), the output of the module are written in the tensordict

[FEATURE] TensorDictModule in_keys allowed as Dict[str, tuple | list] to enable multi use of a sample feature #1101

[FEATURE] TensorDictModule in_keys allowed as Dict[str, tuple | list] to enable multi use of a sample feature #1101

Conversation

bachdj-px commented Nov 21, 2024

Description

Motivation and Context

Types of changes

Checklist

facebook-github-bot commented Nov 21, 2024

Action Required

Process

vmoens left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vmoens left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bachdj-px commented Nov 27, 2024

vmoens left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vmoens commented Nov 29, 2024