Searched refs:padded_base_shape (Results 1 – 2 of 2) sorted by relevance
285 Shape padded_base_shape = base_shape; in GetPaddedShapeForUnevenPartitioning() local286 for (int64 i = 0; i < padded_base_shape.rank(); ++i) { in GetPaddedShapeForUnevenPartitioning()287 padded_base_shape.set_dimensions( in GetPaddedShapeForUnevenPartitioning()290 return padded_base_shape; in GetPaddedShapeForUnevenPartitioning()295 auto padded_base_shape = in PadBaseShapeBeforeUnevenTiledSharding() local297 if (ShapeUtil::Compatible(padded_base_shape, hlo->shape())) { in PadBaseShapeBeforeUnevenTiledSharding()300 return PadToShape(hlo, padded_base_shape, b); in PadBaseShapeBeforeUnevenTiledSharding()
1027 auto padded_base_shape = shard_shape; in ReshardFromPartialReplicateWithDynamicSlice() local1028 for (int64 i = 0; i < padded_base_shape.rank(); ++i) { in ReshardFromPartialReplicateWithDynamicSlice()1029 padded_base_shape.set_dimensions( in ReshardFromPartialReplicateWithDynamicSlice()1030 i, padded_base_shape.dimensions(i) * in ReshardFromPartialReplicateWithDynamicSlice()1033 auto offsets = MakePartitionOffsets(padded_base_shape, temp_target_sharding, in ReshardFromPartialReplicateWithDynamicSlice()1035 auto old_offsets = MakePartitionOffsets(padded_base_shape, sharding(), in ReshardFromPartialReplicateWithDynamicSlice()