Searched refs:tile_shape (Results 1 – 9 of 9) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | sharding_builder.cc | 41 OpSharding Tile(const Shape& tile_shape, in Tile() argument
45 *result.mutable_tile_shape() = tile_shape.ToProto(); in Tile()
55 OpSharding Tile1D(const Shape& tile_shape, int64 num_tiles) { in Tile1D() argument
59 CHECK_EQ(tile_shape.rank(), 1); in Tile1D()
61 *result.mutable_tile_shape() = tile_shape.ToProto(); in Tile1D()
|
| D | sharding_builder.h | 48 OpSharding Tile(const Shape& tile_shape, const TileAssignment& tile_assignment);
54 OpSharding Tile1D(const Shape& tile_shape, int64 num_tiles);
|
| /external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/ |
| D | xla_sharding.py | 436 tile_shape = [
441 tile_shape.append(_np.prod(device_mesh.shape) // _np.prod(tile_shape))
442 tile_assignment = _np.reshape(tile_assignment, tile_shape)
|
| /external/tensorflow/tensorflow/compiler/xla/ |
| D | xla_data.proto | 638 // None of the above; tile_shape and tile_assignment are both used.
646 ShapeProto tile_shape = 2; field
648 // tile_shape and the product of its dimensions must equal
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | sharding_propagation.cc | 1187 const int64 tile_shape = CeilOfRatio(user.shape().dimensions(cdim), in GetShardingFromUser() local
1191 start_indices[cdim] = start_offset / tile_shape; in GetShardingFromUser()
1193 start_offset + instruction.shape().dimensions(cdim), tile_shape); in GetShardingFromUser()
|
| D | hlo_sharding_util.cc | 314 const Shape tile_shape = sharding.TileShape(source_shape); in ReshapeSharding() local
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/ |
| D | tpu_rewrite.mlir | 1893 // tile_shape {
2000 // tile_shape {
2084 // tile_shape {
2168 // tile_shape {
2253 // tile_shape {
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | array_ops.py | 5427 tile_shape = concat(((1,), indices_internal_shape, (1,)), axis=0)
5428 index_grid = tile(index_grid, multiples=tile_shape)
|
| /external/tensorflow/tensorflow/python/keras/ |
| D | backend.py | 4697 tile_shape = array_ops.where_v2(shape_diff > 0, expr_shape,
4699 condition = array_ops.tile(condition, tile_shape)
|