We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Auto wrapping Huggingface models with FullyShardedDataParallel can break them
## 🐛 Bug

If you auto wrap a HuggingFace model with `FullyShardedDataParallel`, depending on what sub-modules are wrapped, and how the output of the modules are accessed in the model, calling `forward` can result in exceptions. 

## What happens

Some sub-modules of the HuggingFace model output a child implementation of [`ModelOutput`](https://github.com/huggingface/transformers/blob/8b240a06617455eae59e1116af6a1a016664e963/src/transformers/file_utils.py#L1884), which is itself a descendent of `OrderedDict`.

An example is, for instance, [`CausalLMOutputWithCrossAttentions` ](https://github.com/huggingface/transformers/blob/8b240a06617455eae59e1116af6a1a016664e963/src/transformers/modeling_outputs.py#L375).

You can access attributes of `ModelOutput` instance derivatives in three ways, e.g. like this:

- `hidden_states = transformer_outputs[0]`
- `hidden_states = transformer_outputs.hidden_states`
- `hidden_states = transformer_outputs["hidden_states"]`

In the HuggingFace library these three types of accessing attributes are used interchangeably.

Now, when a sub-module, which outputs a `ModelOutput` derived instance,  is wrapped with `FullyShardedDataParallel`, this output instance is converted in to an `OrderedDict`, in the [`apply_to_tensors`](https://github.com/facebookresearch/fairscale/blob/fecb665b812b6bfc38442e1fb1557e21508917f4/fairscale/utils/containers.py#L21) function, in `utils/containers.py`, since the `ModelOutput` derives from `OrderedDict`! :

```
        elif isinstance(x, OrderedDict):
            od = OrderedDict()
            for key, value in x.items():
                od[key] = _apply(value)
            return od
```

After it has been converted into a plain `OrderedDict` by `FullyShardedDataParallel`, when the output of the submodule is, for instance, accessed using an index (`transformer_outputs[0]`) or attribute accessor (`transformer_outputs. hidden_states`), an exception is thrown, because the `OrderedDict` instance doesn't know how to handle that. 

## A simple solution

A simple solution is to replace the `OrderedDict` instantiation above with `od = x.__class__()`, thus:

```
        elif isinstance(x, OrderedDict):
            od = x.__class__()
            for key, value in x.items():
                od[key] = _apply(value)
            return od
```

In this way we keep on using the original class, with its extended means to access attributes, which is used in the HuggingFace library.


## To Reproduce

Steps to reproduce the behavior:

1. Use a HuggingFace model, e.g. `GPT2LMHeadModel`
```
from transformers import GPT2Config, GPT2LMHeadModel
```

2. Build the model 
```
        # You could also try, e.g. "gpt2-large"
        model_config = GPT2Config.from_pretrained("gpt2")
        self._model = GPT2LMHeadModel(model_config)

        self._model.to(self._device)
```

3. Auto wrap the model
```
                wrap_policy = functools.partial(default_auto_wrap_policy,
                                                min_num_params=1e8)

                fsdp_params = {
                    "mixed_precision": False,
                    "move_params_to_cpu": False
                }

                with enable_wrap(wrapper_cls=FullyShardedDDP, auto_wrap_policy=wrap_policy, **fsdp_params):
                    self._training_model = auto_wrap(self._model)

                # Finally wrap root module
                self._training_model = FullyShardedDDP(self._training_model, **fsdp_params)
```

4. Use this in a training loop

This results in errors similar to this one where the `OrderedDict` accessed in invalid ways:
```
... transformers/models/gpt2/modeling_gpt2.py", line 1057, in forward
    hidden_states = transformer_outputs[0]
KeyError: 0
```

## Expected behavior

No exception should occur, `FullyShardedDataParallel` should not break calling the HuggingFace model.

## Environment

```
$ pip show fairscale
Name: fairscale
Version: 0.4.3

$ pip show transformers
Name: transformers
Version: 4.12.5

$ pip show torch
Name: torch
Version: 1.10.0

$ python --version
Python 3.7.4
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `fairscale/utils/containers.py`
Content:
```
1 # Copyright (c) Facebook, Inc. and its affiliates.
2 #
3 # This source code is licensed under the BSD license found in the
4 # LICENSE file in the root directory of this source tree.
5 
6 from collections import OrderedDict
7 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
8 
9 import torch
10 from torch.nn.utils.rnn import PackedSequence
11 
12 """Useful functions to deal with tensor types with other python container types."""
13 
14 
15 def apply_to_tensors(fn: Callable, container: Union[torch.Tensor, Dict, List, Tuple, Set]) -> Any:
16     """Recursively apply to all tensor in different kinds of container types."""
17 
18     def _apply(x: Union[torch.Tensor, Dict, List, Tuple, Set]) -> Any:
19         if torch.is_tensor(x):
20             return fn(x)
21         elif isinstance(x, OrderedDict):
22             od = OrderedDict()
23             for key, value in x.items():
24                 od[key] = _apply(value)
25             return od
26         elif isinstance(x, PackedSequence):
27             _apply(x.data)
28             return x
29         elif isinstance(x, dict):
30             return {key: _apply(value) for key, value in x.items()}
31         elif isinstance(x, list):
32             return [_apply(x) for x in x]
33         elif isinstance(x, tuple):
34             return tuple(_apply(x) for x in x)
35         elif isinstance(x, set):
36             return {_apply(x) for x in x}
37         else:
38             return x
39 
40     return _apply(container)
41 
42 
43 def pack_kwargs(*args: Any, **kwargs: Any) -> Tuple[Tuple[str, ...], Tuple[Any, ...]]:
44     """
45     Turn argument list into separate key list and value list (unpack_kwargs does the opposite)
46 
47     Usage::
48 
49         kwarg_keys, flat_args = pack_kwargs(1, 2, a=3, b=4)
50         assert kwarg_keys == ("a", "b")
51         assert flat_args == (1, 2, 3, 4)
52         args, kwargs = unpack_kwargs(kwarg_keys, flat_args)
53         assert args == (1, 2)
54         assert kwargs == {"a": 3, "b": 4}
55     """
56     kwarg_keys: List[str] = []
57     flat_args: List[Any] = list(args)
58     for k, v in kwargs.items():
59         kwarg_keys.append(k)
60         flat_args.append(v)
61     return tuple(kwarg_keys), tuple(flat_args)
62 
63 
64 def unpack_kwargs(kwarg_keys: Tuple[str, ...], flat_args: Tuple[Any, ...]) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
65     """See pack_kwargs."""
66     assert len(kwarg_keys) <= len(flat_args), f"too many keys {len(kwarg_keys)} vs. {len(flat_args)}"
67     if len(kwarg_keys) == 0:
68         return flat_args, {}
69     args = flat_args[: -len(kwarg_keys)]
70     kwargs = {k: v for k, v in zip(kwarg_keys, flat_args[-len(kwarg_keys) :])}
71     return args, kwargs
72 
73 
74 def split_non_tensors(
75     mixed: Union[torch.Tensor, Tuple[Any, ...]]
76 ) -> Tuple[Tuple[torch.Tensor, ...], Optional[Dict[str, List[Any]]]]:
77     """
78     Split a tuple into a list of tensors and the rest with information
79     for later reconstruction.
80 
81     Usage::
82 
83         x = torch.Tensor([1])
84         y = torch.Tensor([2])
85         tensors, packed_non_tensors = split_non_tensors((x, y, None, 3))
86         assert tensors == (x, y)
87         assert packed_non_tensors == {
88             "is_tensor": [True, True, False, False],
89             "objects": [None, 3],
90         }
91         recon = unpack_non_tensors(tensors, packed_non_tensors)
92         assert recon == (x, y, None, 3)
93     """
94     if isinstance(mixed, torch.Tensor):
95         return (mixed,), None
96     tensors: List[torch.Tensor] = []
97     packed_non_tensors: Dict[str, List[Any]] = {"is_tensor": [], "objects": []}
98     for o in mixed:
99         if isinstance(o, torch.Tensor):
100             packed_non_tensors["is_tensor"].append(True)
101             tensors.append(o)
102         else:
103             packed_non_tensors["is_tensor"].append(False)
104             packed_non_tensors["objects"].append(o)
105     return tuple(tensors), packed_non_tensors
106 
107 
108 def unpack_non_tensors(
109     tensors: Tuple[torch.Tensor, ...], packed_non_tensors: Optional[Dict[str, List[Any]]]
110 ) -> Tuple[Any, ...]:
111     """See split_non_tensors."""
112     if packed_non_tensors is None:
113         return tensors
114     assert isinstance(packed_non_tensors, dict), type(packed_non_tensors)
115     mixed: List[Any] = []
116     is_tensor_list = packed_non_tensors["is_tensor"]
117     objects = packed_non_tensors["objects"]
118     assert len(tensors) + len(objects) == len(is_tensor_list), (
119         f"len(tensors) {len(tensors)} len(objects) {len(objects)} " f"len(is_tensor_list) {len(is_tensor_list)}"
120     )
121     obj_i = tnsr_i = 0
122     for is_tensor in is_tensor_list:
123         if is_tensor:
124             mixed.append(tensors[tnsr_i])
125             tnsr_i += 1
126         else:
127             mixed.append(objects[obj_i])
128             obj_i += 1
129     return tuple(mixed)
130
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/fairscale/utils/containers.py b/fairscale/utils/containers.py
--- a/fairscale/utils/containers.py
+++ b/fairscale/utils/containers.py
@@ -19,7 +19,7 @@
         if torch.is_tensor(x):
             return fn(x)
         elif isinstance(x, OrderedDict):
-            od = OrderedDict()
+            od = x.__class__()
             for key, value in x.items():
                 od[key] = _apply(value)
             return od

{"golden_diff": "diff --git a/fairscale/utils/containers.py b/fairscale/utils/containers.py\n--- a/fairscale/utils/containers.py\n+++ b/fairscale/utils/containers.py\n@@ -19,7 +19,7 @@\n         if torch.is_tensor(x):\n             return fn(x)\n         elif isinstance(x, OrderedDict):\n-            od = OrderedDict()\n+            od = x.__class__()\n             for key, value in x.items():\n                 od[key] = _apply(value)\n             return od\n", "issue": "Auto wrapping Huggingface models with FullyShardedDataParallel can break them\n## \ud83d\udc1b Bug\r\n\r\nIf you auto wrap a HuggingFace model with `FullyShardedDataParallel`, depending on what sub-modules are wrapped, and how the output of the modules are accessed in the model, calling `forward` can result in exceptions. \r\n\r\n## What happens\r\n\r\nSome sub-modules of the HuggingFace model output a child implementation of [`ModelOutput`](https://github.com/huggingface/transformers/blob/8b240a06617455eae59e1116af6a1a016664e963/src/transformers/file_utils.py#L1884), which is itself a descendent of `OrderedDict`.\r\n\r\nAn example is, for instance, [`CausalLMOutputWithCrossAttentions` ](https://github.com/huggingface/transformers/blob/8b240a06617455eae59e1116af6a1a016664e963/src/transformers/modeling_outputs.py#L375).\r\n\r\nYou can access attributes of `ModelOutput` instance derivatives in three ways, e.g. like this:\r\n\r\n- `hidden_states = transformer_outputs[0]`\r\n- `hidden_states = transformer_outputs.hidden_states`\r\n- `hidden_states = transformer_outputs[\"hidden_states\"]`\r\n\r\nIn the HuggingFace library these three types of accessing attributes are used interchangeably.\r\n\r\nNow, when a sub-module, which outputs a `ModelOutput` derived instance,  is wrapped with `FullyShardedDataParallel`, this output instance is converted in to an `OrderedDict`, in the [`apply_to_tensors`](https://github.com/facebookresearch/fairscale/blob/fecb665b812b6bfc38442e1fb1557e21508917f4/fairscale/utils/containers.py#L21) function, in `utils/containers.py`, since the `ModelOutput` derives from `OrderedDict`! :\r\n\r\n```\r\n        elif isinstance(x, OrderedDict):\r\n            od = OrderedDict()\r\n            for key, value in x.items():\r\n                od[key] = _apply(value)\r\n            return od\r\n```\r\n\r\nAfter it has been converted into a plain `OrderedDict` by `FullyShardedDataParallel`, when the output of the submodule is, for instance, accessed using an index (`transformer_outputs[0]`) or attribute accessor (`transformer_outputs. hidden_states`), an exception is thrown, because the `OrderedDict` instance doesn't know how to handle that. \r\n\r\n## A simple solution\r\n\r\nA simple solution is to replace the `OrderedDict` instantiation above with `od = x.__class__()`, thus:\r\n\r\n```\r\n        elif isinstance(x, OrderedDict):\r\n            od = x.__class__()\r\n            for key, value in x.items():\r\n                od[key] = _apply(value)\r\n            return od\r\n```\r\n\r\nIn this way we keep on using the original class, with its extended means to access attributes, which is used in the HuggingFace library.\r\n\r\n\r\n## To Reproduce\r\n\r\nSteps to reproduce the behavior:\r\n\r\n1. Use a HuggingFace model, e.g. `GPT2LMHeadModel`\r\n```\r\nfrom transformers import GPT2Config, GPT2LMHeadModel\r\n```\r\n\r\n2. Build the model \r\n```\r\n        # You could also try, e.g. \"gpt2-large\"\r\n        model_config = GPT2Config.from_pretrained(\"gpt2\")\r\n        self._model = GPT2LMHeadModel(model_config)\r\n\r\n        self._model.to(self._device)\r\n```\r\n\r\n3. Auto wrap the model\r\n```\r\n                wrap_policy = functools.partial(default_auto_wrap_policy,\r\n                                                min_num_params=1e8)\r\n\r\n                fsdp_params = {\r\n                    \"mixed_precision\": False,\r\n                    \"move_params_to_cpu\": False\r\n                }\r\n\r\n                with enable_wrap(wrapper_cls=FullyShardedDDP, auto_wrap_policy=wrap_policy, **fsdp_params):\r\n                    self._training_model = auto_wrap(self._model)\r\n\r\n                # Finally wrap root module\r\n                self._training_model = FullyShardedDDP(self._training_model, **fsdp_params)\r\n```\r\n\r\n4. Use this in a training loop\r\n\r\nThis results in errors similar to this one where the `OrderedDict` accessed in invalid ways:\r\n```\r\n... transformers/models/gpt2/modeling_gpt2.py\", line 1057, in forward\r\n    hidden_states = transformer_outputs[0]\r\nKeyError: 0\r\n```\r\n\r\n## Expected behavior\r\n\r\nNo exception should occur, `FullyShardedDataParallel` should not break calling the HuggingFace model.\r\n\r\n## Environment\r\n\r\n```\r\n$ pip show fairscale\r\nName: fairscale\r\nVersion: 0.4.3\r\n\r\n$ pip show transformers\r\nName: transformers\r\nVersion: 4.12.5\r\n\r\n$ pip show torch\r\nName: torch\r\nVersion: 1.10.0\r\n\r\n$ python --version\r\nPython 3.7.4\r\n```\r\n\n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom collections import OrderedDict\nfrom typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union\n\nimport torch\nfrom torch.nn.utils.rnn import PackedSequence\n\n\"\"\"Useful functions to deal with tensor types with other python container types.\"\"\"\n\n\ndef apply_to_tensors(fn: Callable, container: Union[torch.Tensor, Dict, List, Tuple, Set]) -> Any:\n    \"\"\"Recursively apply to all tensor in different kinds of container types.\"\"\"\n\n    def _apply(x: Union[torch.Tensor, Dict, List, Tuple, Set]) -> Any:\n        if torch.is_tensor(x):\n            return fn(x)\n        elif isinstance(x, OrderedDict):\n            od = OrderedDict()\n            for key, value in x.items():\n                od[key] = _apply(value)\n            return od\n        elif isinstance(x, PackedSequence):\n            _apply(x.data)\n            return x\n        elif isinstance(x, dict):\n            return {key: _apply(value) for key, value in x.items()}\n        elif isinstance(x, list):\n            return [_apply(x) for x in x]\n        elif isinstance(x, tuple):\n            return tuple(_apply(x) for x in x)\n        elif isinstance(x, set):\n            return {_apply(x) for x in x}\n        else:\n            return x\n\n    return _apply(container)\n\n\ndef pack_kwargs(*args: Any, **kwargs: Any) -> Tuple[Tuple[str, ...], Tuple[Any, ...]]:\n    \"\"\"\n    Turn argument list into separate key list and value list (unpack_kwargs does the opposite)\n\n    Usage::\n\n        kwarg_keys, flat_args = pack_kwargs(1, 2, a=3, b=4)\n        assert kwarg_keys == (\"a\", \"b\")\n        assert flat_args == (1, 2, 3, 4)\n        args, kwargs = unpack_kwargs(kwarg_keys, flat_args)\n        assert args == (1, 2)\n        assert kwargs == {\"a\": 3, \"b\": 4}\n    \"\"\"\n    kwarg_keys: List[str] = []\n    flat_args: List[Any] = list(args)\n    for k, v in kwargs.items():\n        kwarg_keys.append(k)\n        flat_args.append(v)\n    return tuple(kwarg_keys), tuple(flat_args)\n\n\ndef unpack_kwargs(kwarg_keys: Tuple[str, ...], flat_args: Tuple[Any, ...]) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:\n    \"\"\"See pack_kwargs.\"\"\"\n    assert len(kwarg_keys) <= len(flat_args), f\"too many keys {len(kwarg_keys)} vs. {len(flat_args)}\"\n    if len(kwarg_keys) == 0:\n        return flat_args, {}\n    args = flat_args[: -len(kwarg_keys)]\n    kwargs = {k: v for k, v in zip(kwarg_keys, flat_args[-len(kwarg_keys) :])}\n    return args, kwargs\n\n\ndef split_non_tensors(\n    mixed: Union[torch.Tensor, Tuple[Any, ...]]\n) -> Tuple[Tuple[torch.Tensor, ...], Optional[Dict[str, List[Any]]]]:\n    \"\"\"\n    Split a tuple into a list of tensors and the rest with information\n    for later reconstruction.\n\n    Usage::\n\n        x = torch.Tensor([1])\n        y = torch.Tensor([2])\n        tensors, packed_non_tensors = split_non_tensors((x, y, None, 3))\n        assert tensors == (x, y)\n        assert packed_non_tensors == {\n            \"is_tensor\": [True, True, False, False],\n            \"objects\": [None, 3],\n        }\n        recon = unpack_non_tensors(tensors, packed_non_tensors)\n        assert recon == (x, y, None, 3)\n    \"\"\"\n    if isinstance(mixed, torch.Tensor):\n        return (mixed,), None\n    tensors: List[torch.Tensor] = []\n    packed_non_tensors: Dict[str, List[Any]] = {\"is_tensor\": [], \"objects\": []}\n    for o in mixed:\n        if isinstance(o, torch.Tensor):\n            packed_non_tensors[\"is_tensor\"].append(True)\n            tensors.append(o)\n        else:\n            packed_non_tensors[\"is_tensor\"].append(False)\n            packed_non_tensors[\"objects\"].append(o)\n    return tuple(tensors), packed_non_tensors\n\n\ndef unpack_non_tensors(\n    tensors: Tuple[torch.Tensor, ...], packed_non_tensors: Optional[Dict[str, List[Any]]]\n) -> Tuple[Any, ...]:\n    \"\"\"See split_non_tensors.\"\"\"\n    if packed_non_tensors is None:\n        return tensors\n    assert isinstance(packed_non_tensors, dict), type(packed_non_tensors)\n    mixed: List[Any] = []\n    is_tensor_list = packed_non_tensors[\"is_tensor\"]\n    objects = packed_non_tensors[\"objects\"]\n    assert len(tensors) + len(objects) == len(is_tensor_list), (\n        f\"len(tensors) {len(tensors)} len(objects) {len(objects)} \" f\"len(is_tensor_list) {len(is_tensor_list)}\"\n    )\n    obj_i = tnsr_i = 0\n    for is_tensor in is_tensor_list:\n        if is_tensor:\n            mixed.append(tensors[tnsr_i])\n            tnsr_i += 1\n        else:\n            mixed.append(objects[obj_i])\n            obj_i += 1\n    return tuple(mixed)\n", "path": "fairscale/utils/containers.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom collections import OrderedDict\nfrom typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union\n\nimport torch\nfrom torch.nn.utils.rnn import PackedSequence\n\n\"\"\"Useful functions to deal with tensor types with other python container types.\"\"\"\n\n\ndef apply_to_tensors(fn: Callable, container: Union[torch.Tensor, Dict, List, Tuple, Set]) -> Any:\n    \"\"\"Recursively apply to all tensor in different kinds of container types.\"\"\"\n\n    def _apply(x: Union[torch.Tensor, Dict, List, Tuple, Set]) -> Any:\n        if torch.is_tensor(x):\n            return fn(x)\n        elif isinstance(x, OrderedDict):\n            od = x.__class__()\n            for key, value in x.items():\n                od[key] = _apply(value)\n            return od\n        elif isinstance(x, PackedSequence):\n            _apply(x.data)\n            return x\n        elif isinstance(x, dict):\n            return {key: _apply(value) for key, value in x.items()}\n        elif isinstance(x, list):\n            return [_apply(x) for x in x]\n        elif isinstance(x, tuple):\n            return tuple(_apply(x) for x in x)\n        elif isinstance(x, set):\n            return {_apply(x) for x in x}\n        else:\n            return x\n\n    return _apply(container)\n\n\ndef pack_kwargs(*args: Any, **kwargs: Any) -> Tuple[Tuple[str, ...], Tuple[Any, ...]]:\n    \"\"\"\n    Turn argument list into separate key list and value list (unpack_kwargs does the opposite)\n\n    Usage::\n\n        kwarg_keys, flat_args = pack_kwargs(1, 2, a=3, b=4)\n        assert kwarg_keys == (\"a\", \"b\")\n        assert flat_args == (1, 2, 3, 4)\n        args, kwargs = unpack_kwargs(kwarg_keys, flat_args)\n        assert args == (1, 2)\n        assert kwargs == {\"a\": 3, \"b\": 4}\n    \"\"\"\n    kwarg_keys: List[str] = []\n    flat_args: List[Any] = list(args)\n    for k, v in kwargs.items():\n        kwarg_keys.append(k)\n        flat_args.append(v)\n    return tuple(kwarg_keys), tuple(flat_args)\n\n\ndef unpack_kwargs(kwarg_keys: Tuple[str, ...], flat_args: Tuple[Any, ...]) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:\n    \"\"\"See pack_kwargs.\"\"\"\n    assert len(kwarg_keys) <= len(flat_args), f\"too many keys {len(kwarg_keys)} vs. {len(flat_args)}\"\n    if len(kwarg_keys) == 0:\n        return flat_args, {}\n    args = flat_args[: -len(kwarg_keys)]\n    kwargs = {k: v for k, v in zip(kwarg_keys, flat_args[-len(kwarg_keys) :])}\n    return args, kwargs\n\n\ndef split_non_tensors(\n    mixed: Union[torch.Tensor, Tuple[Any, ...]]\n) -> Tuple[Tuple[torch.Tensor, ...], Optional[Dict[str, List[Any]]]]:\n    \"\"\"\n    Split a tuple into a list of tensors and the rest with information\n    for later reconstruction.\n\n    Usage::\n\n        x = torch.Tensor([1])\n        y = torch.Tensor([2])\n        tensors, packed_non_tensors = split_non_tensors((x, y, None, 3))\n        assert tensors == (x, y)\n        assert packed_non_tensors == {\n            \"is_tensor\": [True, True, False, False],\n            \"objects\": [None, 3],\n        }\n        recon = unpack_non_tensors(tensors, packed_non_tensors)\n        assert recon == (x, y, None, 3)\n    \"\"\"\n    if isinstance(mixed, torch.Tensor):\n        return (mixed,), None\n    tensors: List[torch.Tensor] = []\n    packed_non_tensors: Dict[str, List[Any]] = {\"is_tensor\": [], \"objects\": []}\n    for o in mixed:\n        if isinstance(o, torch.Tensor):\n            packed_non_tensors[\"is_tensor\"].append(True)\n            tensors.append(o)\n        else:\n            packed_non_tensors[\"is_tensor\"].append(False)\n            packed_non_tensors[\"objects\"].append(o)\n    return tuple(tensors), packed_non_tensors\n\n\ndef unpack_non_tensors(\n    tensors: Tuple[torch.Tensor, ...], packed_non_tensors: Optional[Dict[str, List[Any]]]\n) -> Tuple[Any, ...]:\n    \"\"\"See split_non_tensors.\"\"\"\n    if packed_non_tensors is None:\n        return tensors\n    assert isinstance(packed_non_tensors, dict), type(packed_non_tensors)\n    mixed: List[Any] = []\n    is_tensor_list = packed_non_tensors[\"is_tensor\"]\n    objects = packed_non_tensors[\"objects\"]\n    assert len(tensors) + len(objects) == len(is_tensor_list), (\n        f\"len(tensors) {len(tensors)} len(objects) {len(objects)} \" f\"len(is_tensor_list) {len(is_tensor_list)}\"\n    )\n    obj_i = tnsr_i = 0\n    for is_tensor in is_tensor_list:\n        if is_tensor:\n            mixed.append(tensors[tnsr_i])\n            tnsr_i += 1\n        else:\n            mixed.append(objects[obj_i])\n            obj_i += 1\n    return tuple(mixed)\n", "path": "fairscale/utils/containers.py"}]}