mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
23b2fba79a
Summary: Stacked PRs * #33474 - [jit] Remove list specializations from pickler * **#33255 - [jit] Add type tags to lists/dicts in pickle** This adds a global call to `torch.jit._pickle.restore_type_tags` for lists and dicts so that we can preserve their types after serialization. ](https://our.intern.facebook.com/intern/diff/20346780/) Pull Request resolved: https://github.com/pytorch/pytorch/pull/33255 Pulled By: driazati Differential Revision: D20346780 fbshipit-source-id: c8534954ef4adb2e3c880401acbee30cd284f3db
37 lines
954 B
Python
37 lines
954 B
Python
# These functions are referenced from the pickle archives produced by
|
|
# ScriptModule.save()
|
|
|
|
|
|
# These (`build_*`) functions used to be used by `pickler.cpp` to specify
|
|
# the type of the list for certain special types, but now all lists get
|
|
# a type attached and restored via `restore_type_tag` below. The legacy
|
|
# functions should stick around for backwards-compatibility.
|
|
|
|
def build_intlist(data):
|
|
return data
|
|
|
|
|
|
def build_tensorlist(data):
|
|
return data
|
|
|
|
|
|
def build_doublelist(data):
|
|
return data
|
|
|
|
|
|
def build_boollist(data):
|
|
return data
|
|
|
|
|
|
def build_tensor_from_id(data):
|
|
if isinstance(data, int):
|
|
# just the id, can't really do anything
|
|
return data
|
|
|
|
|
|
def restore_type_tag(value, type_str):
|
|
# The type_ptr is used by the jit unpickler to restore the full static type
|
|
# to container types like list when they are re-loaded, but this doesn't
|
|
# matter for Python, so just return the plain value
|
|
return value
|