ttnn.create_sharded_memory_config

ttnn.create_sharded_memory_config(shape: Shape | Tuple[int, ...] | List[int], core_grid: CoreGrid | CoreRangeSet, strategy: ShardStrategy, orientation: ShardOrientation | None = None, use_height_and_width_as_shard_shape: bool = False) → MemoryConfig

Creates a MemoryConfig object with a sharding spec, required for sharded ops.

Parameters:

  • shape (ttnn.Shape | Tuple[int, ...] | List[int])the shape of the tensor.

  • core_grid (ttnn.CoreGrid | ttnn.CoreRangeSet)the core_grid on which to distribute the sharded tensor on (writes to the cores L1s).

  • strategy (ttnn.ShardStrategy)the sharding strategy of either height, width or block.

  • orientation (ttnn.ShardOrientation, optional)the order in which to traverse the cores when reading/writing shards. Defaults to None.

  • use_height_and_width_as_shard_shape (bool, optional)if True, the height and width of the tensor will be used as the shard shape. Defaults to False. If is False, the shard shape will be calculated based on the core_grid and the tensor shape where tensor shape is seen as [math.prod(dims), width]

Returns:

ttnn.MemoryConfigthe MemoryConfig object.

Note

Currently sharding only supports L1 tensors.

Example

>>> tensor = ttnn.create_sharded_memory_config((5, 8), (320,64), ttnn.ShardStrategy.BLOCK, ttnn.ShardOrientation.ROW_MAJOR)