Warning

The named tensor API is a prototype feature and subject to change.

Creating named tensors#

Factory functions now take a newnames argument that associates a namewith each dimension.

>>>torch.zeros(2,3,names=('N','C'))tensor([[0.,0.,0.],[0.,0.,0.]],names=('N','C'))

Named dimensions, like regular Tensor dimensions, are ordered.tensor.names[i] is the name of dimensioni oftensor.

The following factory functions support named tensors:

• torch.empty()

• torch.rand()

• torch.randn()

• torch.ones()

• torch.tensor()

• torch.zeros()

Named dimensions#

Seenames for restrictions on tensor names.

Usenames to access the dimension names of a tensor andrename() to rename named dimensions.

>>>imgs=torch.randn(1,2,2,3,names=('N','C','H','W'))>>>imgs.names('N','C','H','W')>>>renamed_imgs=imgs.rename(H='height',W='width')>>>renamed_imgs.names('N','C','height','width)

Named tensors can coexist with unnamed tensors; named tensors are instances oftorch.Tensor. Unnamed tensors haveNone-named dimensions. Namedtensors do not require all dimensions to be named.

>>>imgs=torch.randn(1,2,2,3,names=(None,'C','H','W'))>>>imgs.names(None,'C','H','W')

Name propagation semantics#

Named tensors use names to automatically check that APIs are being calledcorrectly at runtime. This occurs in a process calledname inference.More formally, name inference consists of the following two steps:

• Check names: an operator may perform automatic checks at runtime thatcheck that certain dimension names must match.

• Propagate names: name inference propagates names to output tensors.

All operations that support named tensors propagate names.

>>>x=torch.randn(3,3,names=('N','C'))>>>x.abs().names('N','C')

x=torch.randn(3,names=('X',))y=torch.randn(3)z=torch.randn(3,names=('Z',))

>>># x + y  # match('X', None) is True>>># x + z  # match('X', 'Z') is False>>># x + x  # match('X', 'X') is True>>>x+zErrorwhenattemptingtobroadcastdims['X']anddims['Z']:dim'X'anddim'Z'areatthesamepositionfromtherightbutdonotmatch.

>>>(x+y).names('X',)>>>(x+x).names('X',)

Explicit alignment by names#

Usealign_as() oralign_to() to align tensor dimensionsby name to a specified ordering. This is useful for performing “broadcasting by names”.

# This function is agnostic to the dimension ordering of `input`,# as long as it has a `C` dimension somewhere.defscale_channels(input,scale):scale=scale.refine_names('C')returninput*scale.align_as(input)>>>num_channels=3>>>scale=torch.randn(num_channels,names=('C',))>>>imgs=torch.rand(3,3,3,num_channels,names=('N','H','W','C'))>>>more_imgs=torch.rand(3,num_channels,3,3,names=('N','C','H','W'))>>>videos=torch.randn(3,num_channels,3,3,3,names=('N','C','H','W','D')>>>scale_channels(imgs,scale)>>>scale_channels(more_imgs,scale)>>>scale_channels(videos,scale)

Manipulating dimensions#

Usealign_to() to permute large amounts of dimensions withoutmentioning all of them as in required bypermute().

>>>tensor=torch.randn(2,2,2,2,2,2)>>>named_tensor=tensor.refine_names('A','B','C','D','E','F')# Move the F (dim 5) and E dimension (dim 4) to the front while keeping# the rest in the same order>>>tensor.permute(5,4,0,1,2,3)>>>named_tensor.align_to('F','E',...)

Useflatten() andunflatten() to flatten and unflattendimensions, respectively. These methods are more verbose thanview()andreshape(), but have more semantic meaning to someone reading the code.

>>>imgs=torch.randn(32,3,128,128)>>>named_imgs=imgs.refine_names('N','C','H','W')>>>flat_imgs=imgs.view(32,-1)>>>named_flat_imgs=named_imgs.flatten(['C','H','W'],'features')>>>named_flat_imgs.names('N','features')>>>unflattened_named_imgs=named_flat_imgs.unflatten('features',[('C',3),('H',128),('W',128)])>>>unflattened_named_imgs.names('N','C','H','W')

Autograd support#

Autograd currently supports named tensors in a limited manner: autograd ignoresnames on all tensors. Gradient computation is still correct but we lose thesafety that names give us.

>>>x=torch.randn(3,names=('D',))>>>weight=torch.randn(3,names=('D',),requires_grad=True)>>>loss=(x-weight).abs()>>>grad_loss=torch.randn(3)>>>loss.backward(grad_loss)>>>weight.grad# Unnamed for now. Will be named in the futuretensor([-1.8107,-0.6357,0.0783])>>>weight.grad.zero_()>>>grad_loss=grad_loss.refine_names('C')>>>loss=(x-weight).abs()# Ideally we'd check that the names of loss and grad_loss match but we don't yet.>>>loss.backward(grad_loss)>>>weight.gradtensor([-1.8107,-0.6357,0.0783])

Currently supported operations and subsystems#

Named tensor API reference#

In this section please find the documentation for named tensor specific APIs.For a comprehensive reference for how names are propagated through other PyTorchoperators, seeNamed Tensors operator coverage.

classtorch.Tensor

names#

Stores names for each of this tensor’s dimensions.

names[idx] corresponds to the name of tensor dimensionidx.Names are either a string if the dimension is named orNone if thedimension is unnamed.

Dimension names may contain characters or underscore. Furthermore, a dimensionname must be a valid Python variable name (i.e., does not start with underscore).

Tensors may not have two named dimensions with the same name.

Warning

The named tensor API is experimental and subject to change.

rename(*names,**rename_map)[source]#

Renames dimension names ofself.

There are two main usages:

self.rename(**rename_map) returns a view on tensor that has dimsrenamed as specified in the mappingrename_map.

self.rename(*names) returns a view on tensor, renaming alldimensions positionally usingnames.Useself.rename(None) to drop names on a tensor.

One cannot specify both positional argsnames and keyword argsrename_map.

Examples:

>>>imgs=torch.rand(2,3,5,7,names=('N','C','H','W'))>>>renamed_imgs=imgs.rename(N='batch',C='channels')>>>renamed_imgs.names('batch', 'channels', 'H', 'W')>>>renamed_imgs=imgs.rename(None)>>>renamed_imgs.names(None, None, None, None)>>>renamed_imgs=imgs.rename('batch','channel','height','width')>>>renamed_imgs.names('batch', 'channel', 'height', 'width')

Warning

The named tensor API is experimental and subject to change.

rename_(*names,**rename_map)[source]#

In-place version ofrename().

refine_names(*names)[source]#

Refines the dimension names ofself according tonames.

Refining is a special case of renaming that “lifts” unnamed dimensions.ANone dim can be refined to have any name; a named dim can only berefined to have the same name.

Because named tensors can coexist with unnamed tensors, refining namesgives a nice way to write named-tensor-aware code that works with bothnamed and unnamed tensors.

names may contain up to one Ellipsis (...).The Ellipsis is expanded greedily; it is expanded in-place to fillnames to the same length asself.dim() using names from thecorresponding indices ofself.names.

Python 2 does not support Ellipsis but one may use a string literalinstead ('...').

Parameters

names (iterable ofstr) – The desired names of the output tensor. Maycontain up to one Ellipsis.

Examples:

>>>imgs=torch.randn(32,3,128,128)>>>named_imgs=imgs.refine_names('N','C','H','W')>>>named_imgs.names('N', 'C', 'H', 'W')>>>tensor=torch.randn(2,3,5,7,11)>>>tensor=tensor.refine_names('A',...,'B','C')>>>tensor.names('A', None, None, 'B', 'C')

Warning

The named tensor API is experimental and subject to change.

align_as(other)→Tensor#

Permutes the dimensions of theself tensor to match the dimension orderin theother tensor, adding size-one dims for any new names.

This operation is useful for explicit broadcasting by names (see examples).

All of the dims ofself must be named in order to use this method.The resulting tensor is a view on the original tensor.

All dimension names ofself must be present inother.names.other may contain named dimensions that are not inself.names;the output tensor has a size-one dimension for each of those new names.

To align a tensor to a specific order, usealign_to().

Examples:

# Example 1: Applying a mask>>>mask=torch.randint(2,[127,128],dtype=torch.bool).refine_names('W','H')>>>imgs=torch.randn(32,128,127,3,names=('N','H','W','C'))>>>imgs.masked_fill_(mask.align_as(imgs),0)# Example 2: Applying a per-channel-scale>>>defscale_channels(input,scale):>>>scale=scale.refine_names('C')>>>returninput*scale.align_as(input)>>>num_channels=3>>>scale=torch.randn(num_channels,names=('C',))>>>imgs=torch.rand(32,128,128,num_channels,names=('N','H','W','C'))>>>more_imgs=torch.rand(32,num_channels,128,128,names=('N','C','H','W'))>>>videos=torch.randn(3,num_channels,128,128,128,names=('N','C','H','W','D'))# scale_channels is agnostic to the dimension order of the input>>>scale_channels(imgs,scale)>>>scale_channels(more_imgs,scale)>>>scale_channels(videos,scale)

Warning

The named tensor API is experimental and subject to change.

align_to(*names)[source]#

Permutes the dimensions of theself tensor to match the orderspecified innames, adding size-one dims for any new names.

All of the dims ofself must be named in order to use this method.The resulting tensor is a view on the original tensor.

All dimension names ofself must be present innames.names may contain additional names that are not inself.names;the output tensor has a size-one dimension for each of those new names.

names may contain up to one Ellipsis (...).The Ellipsis is expanded to be equal to all dimension names ofselfthat are not mentioned innames, in the order that they appearinself.

Python 2 does not support Ellipsis but one may use a string literalinstead ('...').

Parameters

names (iterable ofstr) – The desired dimension ordering of theoutput tensor. May contain up to one Ellipsis that is expandedto all unmentioned dim names ofself.

Examples:

>>>tensor=torch.randn(2,2,2,2,2,2)>>>named_tensor=tensor.refine_names('A','B','C','D','E','F')# Move the F and E dims to the front while keeping the rest in order>>>named_tensor.align_to('F','E',...)

Warning

The named tensor API is experimental and subject to change.

flatten(dims,out_dim)→Tensor

Flattensdims into a single dimension with nameout_dim.

All ofdims must be consecutive in order in theself tensor,but not necessary contiguous in memory.

Examples:

>>>imgs=torch.randn(32,3,128,128,names=('N','C','H','W'))>>>flat_imgs=imgs.flatten(['C','H','W'],'features')>>>flat_imgs.names,flat_imgs.shape(('N', 'features'), torch.Size([32, 49152]))

Warning

The named tensor API is experimental and subject to change.

• torchao
• torchrec
• torchft
• TorchCodec
• torchvision
• ExecuTorch
• PyTorch on XLA Devices