Skip to content

Commit 070eb50

Browse files
lucylqfacebook-github-bot
authored andcommitted
Get tensor data
Summary: Retrieve tensor data as a buffer (like executorch_program.buffer) Differential Revision: D82061884
1 parent 8973eeb commit 070eb50

File tree

1 file changed

+17
-0
lines changed

1 file changed

+17
-0
lines changed

exir/program/_program.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -626,6 +626,23 @@ def buffer(self) -> bytes:
626626
self._buffer = bytes(self._get_pte_data())
627627
return self._buffer
628628

629+
@property
630+
def tensor_data(self) -> Dict[str, bytes]:
631+
"""Returns the tensor data as a dictionary of filename to byte data.
632+
633+
Returns:
634+
Dict[str, bytes]: Dictionary mapping tensor data filenames to
635+
their byte content.
636+
Returns empty dict if no tensor data is available.
637+
"""
638+
if self._pte_data is None:
639+
self._get_pte_data() # This populates _tensor_data
640+
641+
if self._tensor_data is None:
642+
return {}
643+
644+
return {filename: bytes(cord) for filename, cord in self._tensor_data.items()}
645+
629646
@property
630647
def program(self) -> Program:
631648
return self._get_emitter_output().program

0 commit comments

Comments
 (0)