Skip to content

torch_to_nnef.tensor.named

Classes:

Name Description
NamedTensor

Tensor enriched with name attribute.

Functions:

Name Description
apply_name_to_tensor_in_module

Transform torch.Tensor or Parameters into NamedTensor.

find_fake_mode

Return the FakeTensorMode of any fake operand, else None.

NamedTensor

NamedTensor(fp_tensor: Tensor, nnef_name: str)

Bases: Tensor

Tensor enriched with name attribute.

Attributes:

Name Type Description
data

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

data property writable

data

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

apply_name_to_tensor_in_module

apply_name_to_tensor_in_module(model: Module)

Transform torch.Tensor or Parameters into NamedTensor.

This is applied at export time of torch_to_nnef Just before doing any tracing and allow to keep variable naming identical to PyTorch one

This consistent naming unlock subsequent manipulations such as LORA applications @ inference or such.

find_fake_mode

find_fake_mode(*containers)

Return the FakeTensorMode of any fake operand, else None.

Opaque float weights are traced as FakeTensor placeholders (see torch_to_nnef.tensor.opaque._fake_trace_tensor). A fake tensor's dispatch only works while its owning mode is active, so when such a tensor meets a :class:NamedTensor (which drops out of the subclass dispatch) the bare op raises TypeError: unsupported operand. We re-activate the mode around the op to let the two compose.