ttnn.as_tensor

ttnn.as_tensor(tensor: torch.Tensor, dtype: ttnn.DataType | None, *, layout: ttnn.Layout = ttnn.ROW_MAJOR_LAYOUT, device: ttnn.MeshDevice = None, memory_config: ttnn.MemoryConfig = None, cache_file_name: str | pathlib.Path = None, preprocess: Callable[[ttnn.Tensor], ttnn.Tensor] = None, mesh_mapper: ttnn.CppTensorToMesh = None) → ttnn.Tensor

Converts the torch.Tensor tensor into a ttnn.Tensor.

Parameters:

  • tensor (torch.Tensor) – the input tensor.

  • dtype (ttnn.DataType, optional) – The ttnn data type.

Keyword Arguments:

  • layout (ttnn.Layout, optional) – The ttnn layout. Defaults to ttnn.ROW_MAJOR_LAYOUT.

  • device (ttnn.MeshDevice, optional) – The ttnn device. Defaults to None.

  • memory_config (ttnn.MemoryConfig, optional) – The ttnn memory configuration. Defaults to None.

  • cache_file_name (str | pathlib.Path, optional) – The cache file name. Defaults to None.

  • preprocess (Callable[[ttnn.Tensor], ttnn.Tensor], optional) – The function to preprocess the tensor before serializing/converting to ttnn. Defaults to None.

  • mesh_mapper (ttnn.CppTensorToMesh, optional) –

    The TensorToMesh to define the mapping from torch to multi-device. Defaults to None.

    • For Grayskull, the on-device tilizer will truncate mantissa bits for bfp* formats.

    • For Wormhole, the on-device tilizer will raise a runtime error (RTE) for bfp8 but will truncate for bfp4/2 formats.

Returns:

ttnn.Tensor – The resulting ttnn tensor.

Example

# Create a Torch tensor and convert it to a TT-NN tensor
tensor = ttnn.as_tensor(torch.randn((2, 3)), dtype=ttnn.bfloat16)
logger.info(tensor.shape)