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.Tensorthe 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(f"TT-NN tensor shape after changing memory configuration: {ttnn_tensor_memory_config_changed.shape}")