ttnn.to_memory_config

ttnn.to_memory_config(tensor: ttnn.Tensor, memory_config: ttnn.MemoryConfig, dtype: ttnn.DataType = None, output_tensor: ttnn.Tensor = None) → ttnn.Tensor

Converts a tensor to the desired memory configuration. Used for converting tensors to sharded tensors, interleaved tensors, or converting between DRAM and L1 memory.

Parameters:

  • tensor (ttnn.Tensor) – the input tensor to be converted.

  • memory_config (ttnn.MemoryConfig) – the desired memory configuration for the tensor.

  • dtype (ttnn.DataType, optional) – the optional ttnn data type. Defaults to None.

  • output_tensor (ttnn.Tensor, optional) – the optional output tensor. Defaults to None.

Returns:

ttnn.Tensor – the converted tensor.

Example

# Create a TT-NN tensor and change its memory configuration
ttnn_tensor = ttnn.rand((2, 3), dtype=ttnn.bfloat16, device=device)
ttnn_tensor_memory_config_changed = ttnn.to_memory_config(ttnn_tensor, memory_config=ttnn.L1_MEMORY_CONFIG)

logger.info("TT-NN tensor shape after changing memory configuration", ttnn_tensor_memory_config_changed.shape)