def test_copy_it(self): d = immutable_dict([(3, 4), (5, 6)]) l = immutable_list([(3, 4), (5, 6)]) self.assertEqual(d, deepcopy(d)) self.assertEqual(l, deepcopy(l))
from .named_members_polyfill import _named_parameters, _named_buffers from typing import Callable, List, Dict, Any, Tuple, Optional from functools import wraps try: from torchdynamo import disable as disable_torchdynamo except ImportError: def disable_torchdynamo(x): return x pytree._register_pytree_node( immutable_collections.immutable_list, lambda x: (list(x), None), lambda x, c: immutable_collections.immutable_list(x), ) pytree._register_pytree_node( immutable_collections.immutable_dict, lambda x: (list(x.values()), list(x.keys())), lambda x, c: immutable_collections.immutable_dict( {key: value for key, value in zip(c, x)}), ) # TODO - move this to PyTorch core. This overrides the pytree implementation for # dict to maintain parity with Deepmind pytree. Context = Any def _dict_flatten(d: Dict[Any, Any]) -> Tuple[List[Any], Context]: