Source code for torch_geometric.utils._one_hot

from typing import Optional

import torch
from torch import Tensor

[docs]def one_hot( index: Tensor, num_classes: Optional[int] = None, dtype: Optional[torch.dtype] = None, ) -> Tensor: r"""Taskes a one-dimensional :obj:`index` tensor and returns a one-hot encoded representation of it with shape :obj:`[*, num_classes]` that has zeros everywhere except where the index of last dimension matches the corresponding value of the input tensor, in which case it will be :obj:`1`. .. note:: This is a more memory-efficient version of :meth:`torch.nn.functional.one_hot` as you can customize the output :obj:`dtype`. Args: index (torch.Tensor): The one-dimensional input tensor. num_classes (int, optional): The total number of classes. If set to :obj:`None`, the number of classes will be inferred as one greater than the largest class value in the input tensor. (default: :obj:`None`) dtype (torch.dtype, optional): The :obj:`dtype` of the output tensor. """ if index.dim() != 1: raise ValueError("'index' tensor needs to be one-dimensional") if num_classes is None: num_classes = int(index.max()) + 1 out = torch.zeros((index.size(0), num_classes), dtype=dtype, device=index.device) return out.scatter_(1, index.unsqueeze(1), 1)