torch_to_nnef.tensor.opaque

Classes:

Name	Description
`OpaqueTensor`
`OpaqueTensorRef`	Allow to pass through 'tracing'.
`SupportsOffloadState`	Mixin for opaque tensors carrying a compact, picklable offload state.

Functions:

Name	Description
`find_opaque_ref_by_py_id`	Allow to fetch back the opaque parameter once passed the jit 'wall'.
`opaque_to_final_tensor`	Even if OpaqueTensor are composed it exposes fully expanded tensor.
`resolve_offload_state`	Return the class able to rebuild `state`, or `None`.
`set_opaque_tensor_in_params_as_ref`	Transform OpaqueTensor Parameters into OpaqueTensorRef.
`trace_tensor_device_for_func`	Which device an opaque parameter can be traced on for this op.

OpaqueTensor

Bases: Tensor

Methods:

Name	Description
`to_base_tensor`	Wrap _to_base_tensor with jit export infos.
`to_trace_tensor`	Trace an opaque tensor without materializing its backing data.

Attributes:

Name	Type	Description
`data`		Very important to keep access to all special attr of OpaqueTensor.

data `property` `writable`

data

Very important to keep access to all special attr of OpaqueTensor.

to_base_tensor

to_base_tensor()

Wrap _to_base_tensor with jit export infos.

to_trace_tensor

to_trace_tensor(device: str)

Trace an opaque tensor without materializing its backing data.

OpaqueTensorRef

OpaqueTensorRef(meta_tensor: Tensor, opaque_tensor: OpaqueTensor)

Bases: Tensor

Allow to pass through 'tracing'.

SupportsOffloadState

Mixin for opaque tensors carrying a compact, picklable offload state.

A tensor subclass that would otherwise pickle a large fp-shaped payload can instead serialize a small self-describing dict. It declares a unique OFFLOAD_STATE_TAG and implements the three hooks below; registration is automatic so OffloadedTensor dispatches on the serialized state rather than knowing any concrete class.

Methods:

Name	Description
`from_offload_state`	Rebuild the tensor from `state`, placed on `target_device`.

from_offload_state `abstractmethod` `classmethod`

from_offload_state(state, target_device=None)

Rebuild the tensor from state, placed on target_device.

The concrete class owns device placement (None means leave it wherever the state deserializes), so the offload layer needs no device-move API on the reconstructed object.

find_opaque_ref_by_py_id

find_opaque_ref_by_py_id(module: Module, py_id: int)

Allow to fetch back the opaque parameter once passed the jit 'wall'.

opaque_to_final_tensor

opaque_to_final_tensor(rtensor: Tensor) -> torch.Tensor

Even if OpaqueTensor are composed it exposes fully expanded tensor.

So for example: an OffloadedTensor that contains a QTensor will 'load' then 'decompress' to show final fp tensor.

resolve_offload_state

resolve_offload_state(state) -> T.Optional[type]

Return the class able to rebuild state, or None.

None means state is not a compact offload-state payload (e.g. a plain pickled tensor), so callers fall back to legacy handling.

set_opaque_tensor_in_params_as_ref

set_opaque_tensor_in_params_as_ref(model: Module)

Transform OpaqueTensor Parameters into OpaqueTensorRef.

This is applied at export time of torch_to_nnef Just before doing any tracing

trace_tensor_device_for_func

trace_tensor_device_for_func(func) -> T.Optional[str]

Which device an opaque parameter can be traced on for this op.

"meta" keeps shape/dtype without materializing the backing data; "cpu" forces real decompressed values; None falls back to to_base_tensor() (also real values).

Note: a "meta" op only carries shape/dtype, so its result is a meta tensor with no values and no real device. Two shapes of forward are thus unsupported for meta ops (they raise during trace instead of exporting):

reading concrete parameter values (e.g. branching on weight.view(...).argmax());
combining the meta result with a real tensor through a non-meta op (e.g. weight.view(...) * cpu_buffer), which raises a device mismatch.

Only shape-propagating uses that flow into a symbolic op (linear/matmul/select chains) are supported.

The "meta" names below must correspond to genuine view/shape ops whose aten kind lives in ir_op.DERIVED_MODULE_ATTR_OPS (that is where the meta result is recognized as aliasing its constant input); keep the two lists in sync when adding a new view op.