Skip to content

tensor_build

torch_to_nnef.op.aten.tensor_build

affine_grid_generator 

affine_grid_generator(g, node, name_to_tensor, op_helper, inference_target, **kwargs)

Map PyTorch: 'aten:affine_grid_generator' to NNEF.

affine_grid_generator(theta, size, align_corners) builds a sampling grid for F.grid_sample. The base grid -- a fixed set of normalized coordinates in [-1, 1] (or pixel-centered for align_corners=False) -- depends only on the output spatial shape, so we bake it as a constant tensor at trace time. The only runtime cost is a single matmul against theta.

Currently 2-D only (theta shape (N, 2, 3), output (N, H, W, 2)). 3-D affine_grid would need a (D*H*W, 4) base grid + reshape to (N, D, H, W, 3).

arange 

arange(g, node, name_to_tensor, inference_target, **kwargs)

This operator can not be exactly exported to NNEF.

In general NNEF spec is against dynamism it could provide so

we implement it as a simple constant variable.

blackman_window 

blackman_window(g, node, name_to_tensor, **kwargs)

Map PyTorch: 'aten:blackman_window' to a NNEF constant tensor.

copy 

copy(g, node, name_to_tensor, inference_target, torch_graph, null_ref, **kwargs)

Map PyTorch: 'aten:copy', 'aten:clone' to NNEF.

empty_like 

empty_like(**kwargs)

Operator can not be exactly exported to NNEF if dynamic.

With tract we use use expansion

eye 

eye(g, node, name_to_tensor, op_helper, **kwargs)

Map PyTorch: 'aten:eye' to NNEF via the eye fragment.

Builds the identity at runtime from index ranges + broadcast-eq, so both static n (Python int from the trace) and dynamic n (a TensorVariable produced by e.g. aten::size) work without baking an n*n constant into the graph. This is critical for attention-mask construction in LLMs where n = seq_len is a dynamic axis.

fill 

fill(g, node, name_to_tensor, torch_graph, inference_target, op_helper, **kwargs)

Map PyTorch: 'aten:fill', 'aten:fill_' to NNEF.

full 

full(g, node, name_to_tensor, torch_graph, inference_target, **kwargs)

Map PyTorch: 'aten:full' to NNEF.

full_like 

full_like(**kwargs)

Operator can not be exactly exported to NNEF if dynamic.

With tract we use use expansion

hamming_window 

hamming_window(g, node, name_to_tensor, **kwargs)

Map PyTorch: 'aten:hamming_window' to a NNEF constant tensor.

hann_window 

hann_window(g, node, name_to_tensor, **kwargs)

Map PyTorch: 'aten:hann_window' to a NNEF constant tensor.

kaiser_window 

kaiser_window(g, node, name_to_tensor, **kwargs)

Map PyTorch: 'aten:kaiser_window' to a NNEF constant tensor.

Reads an optional beta (default 12.0) at input index 2 when the full positional form is used.

linspace 

linspace(g, node, name_to_tensor, **kwargs)

Map PyTorch: 'aten:linspace' to a NNEF constant tensor.

aten::linspace(start, end, steps, dtype, layout, device, pin_memory) is fully determined at trace time when start, end, steps are static, which is the common case. Bake the result via torch.linspace and register as a constant.

new_empty 

new_empty(g, node, name_to_tensor, torch_graph, inference_target, **kwargs)

Map PyTorch: 'aten:new_empty' to NNEF.

Materialises as zeros (NNEF has no "uninitialized" tensor and real callers immediately fill the buffer); shape comes from size, dtype defaults to the source tensor's dtype.

new_full 

new_full(g, node, name_to_tensor, torch_graph, inference_target, **kwargs)

Map PyTorch: 'aten:new_full' to NNEF.

Layout: (input, size, fill_value, dtype, layout, device, pin_memory). The fill value is folded at trace time into a custom build fn so the constant materialises with the right value.

new_ones 

new_ones(g, node, name_to_tensor, torch_graph, inference_target, **kwargs)

Map PyTorch: 'aten:new_ones' to NNEF.

new_zeros 

new_zeros(g, node, name_to_tensor, torch_graph, inference_target, **kwargs)

Map PyTorch: 'aten:new_zeros' to NNEF.

one_hot 

one_hot(g, node, name_to_tensor, op_helper, inference_target, **kwargs)

Map PyTorch: 'aten:one_hot' to NNEF.

Two paths:

  • TractNNEF: emit tract_core_one_hot(input, axis, dim) and cast to int64. Tract has the op natively (core/src/ops/array/one_hot.rs) so this is the fastest path and produces the smallest graph.
  • Pure NNEF: emit the one_hot fragment which decomposes to eq(unsqueeze(input, axis), classes) -> select using stdlib ops. The classes constant is baked at trace time, pre-reshaped to (1,) * input.rank + (num_classes,) so NNEF eq's exact-rank broadcast rule is satisfied.

Torch's one_hot(input, num_classes) appends the one-hot dim as the trailing axis, so axis = input.rank (the new last position in the rank-(R+1) output) regardless of which path we take.

ones 

ones(g, node, name_to_tensor, torch_graph, inference_target, **kwargs)

This operator can not be exactly exported to NNEF.

In general NNEF spec is against dynamism it could provide so

we implement it as a simple constant variable.

ones_like 

ones_like(**kwargs)

Operator can not be exactly exported to NNEF if dynamic.

With tract we use use expansion

tril 

tril(g, node, name_to_tensor, inference_target, **kwargs)

Map PyTorch: 'aten:tril' to NNEF.

triu 

triu(g, node, name_to_tensor, inference_target, **kwargs)

Map PyTorch: 'aten:triu' to NNEF.

zero 

zero(**kwargs)

Map PyTorch: 'aten:zero' (and zero_) to NNEF.

Tensor.zero_() writes 0 into every position regardless of the original value (even NaN / +/-Inf). Reuse the _x_like machinery that already powers zeros_like so we materialise a true constant of zeros matching the input's shape and dtype: correct on every input, and shares the same dynamic-axes path (tract_core_shape_of + tile expansion) when shapes aren't known at trace time.

Earlier this was implemented as sub(x, x); that produced 0 for finite inputs but NaN for NaN inputs (since NaN - NaN == NaN), which silently diverged from zero_'s set-everything-to-0 semantics.

zeros 

zeros(g, node, name_to_tensor, torch_graph, inference_target, **kwargs)

Map PyTorch: 'aten:zeros' to NNEF.

zeros_like 

zeros_like(**kwargs)

Operator can not be exactly exported to NNEF if dynamic.

With tract we use use exapnsion