ttnn.create_sharded_memory_config

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

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

Args:

shape (ttnn.Shape | Tuple[int, …] | List[int]): the shape of the tensor. core_grid (ttnn.CoreGrid | ttnn.CoreRange): 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. halo (bool, optional): if the shards have overlapping values. Defaults to False. 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.MemoryConfig: the 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, False)