Arbitrary N-D sharding via TensorSpec::sharded(NdShardSpec, ShardShapeAlignment).
Pad the logical tensor, align the shard shape, cut an N-D page grid into
N-D shards (no flatten), then distribute shards to cores
(round_robin or grid_2d). Default demo: logical
[2,8,8] TILE 2×2, shard [2,4,4] on a
2×2 core grid → 4 shards of page-shape [2,2,2] (3D cubes).
2,8,8 or 2,64,64NdShardSpec.grid).ShardDistributionStrategy — round-robin stacks on the full core list; grid_2d needs shard rank ≤ 2 and maps the folded 2D shard grid onto cores.
TensorSpec::ShardShapeAlignment — rounds the trailing shard dims up to the layout alignment.
Alignment{} → create_default_alignment (pads the tensor, not the shard).