Searched refs:shard_shape (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/core/tpu/kernels/xla/ |
| D | infeed_op.cc | 58 auto shard_shape = in UpdateInfeedLayout() local
61 *shape->mutable_layout() = shard_shape.layout(); in UpdateInfeedLayout()
63 if (xla::ShapeUtil::ElementsIn(shard_shape) == 0) { in UpdateInfeedLayout()
68 if (!xla::LayoutUtil::Equal(shard_shape.layout(), shape->layout())) { in UpdateInfeedLayout()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 470 auto shard_shape = MakePartitionedShape(shape, target); in ReshardNoCache() local
472 shard_shape, padded_hlo, in ReshardNoCache()
474 shard_shape.dimensions())); in ReshardNoCache()
568 auto shard_shape = base_shape_; in ReshardAsWindowedInput() local
638 shard_shape.set_dimensions(i, max_shard_size); in ReshardAsWindowedInput()
761 shard_shape, padded_hlo, offsets_on_padded_shape, in ReshardAsWindowedInput()
762 shard_shape.dimensions())); in ReshardAsWindowedInput()
807 padded_shape.dimensions(dim), shard_shape.dimensions(dim), dim, target, in ReshardAsWindowedInput()
868 const Shape& shard_shape = hlo()->shape(); in ReplicatePartial() local
869 Shape target_shape = shard_shape; in ReplicatePartial()
[all …]
|
| D | convolution_handler.cc | 803 auto shard_shape = MakePartitionedShape(output_base_shape, output_sharding); in PartitionConvolutionTiledOutput() local
806 CHECK(ShapeUtil::Compatible(shard_shape, sharded_conv->shape())); in PartitionConvolutionTiledOutput()
810 shard_shape, sharded_conv, in PartitionConvolutionTiledOutput()
812 shard_shape.dimensions())); in PartitionConvolutionTiledOutput()
|
| D | spmd_partitioner_util.cc | 214 auto shard_shape = MakePartitionedShape(shape, sharding); in MakePartitionOffsets() local
218 offset_arrays[i][device] = indices[i] * shard_shape.dimensions(i); in MakePartitionOffsets()
284 auto shard_shape = MakePartitionedShape(base_shape, sharding); in GetPaddedShapeForUnevenPartitioning() local
288 i, shard_shape.dimensions(i) * sharding.tile_assignment().dim(i)); in GetPaddedShapeForUnevenPartitioning()
1726 auto shard_shape = in PerGroupSliceFromReplicated() local
1729 shard_shape, padded_hlo, in PerGroupSliceFromReplicated()
1732 shard_shape.dimensions())); in PerGroupSliceFromReplicated()
|
| /external/tensorflow/tensorflow/python/kernel_tests/ |
| D | embedding_ops_test.py | 151 shard_shape = [vocab_size // num_shards] + shape
153 shard_shape[0] += 1
161 1.0, shape=shard_shape, dtype=dtype, name=param_name)
164 val = (np.random.rand(*shard_shape).astype(np_type)) + 1
|