Searched refs:TileOffsetForDevice (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_sharding_test.cc | 140 EXPECT_EQ(sharding.TileOffsetForDevice(shape, 0), in TEST_F()
142 EXPECT_EQ(sharding.TileOffsetForDevice(shape, 3), in TEST_F()
144 EXPECT_EQ(sharding.TileOffsetForDevice(shape, 2), in TEST_F()
146 EXPECT_EQ(sharding.TileOffsetForDevice(shape, 1), in TEST_F()
|
| D | hlo_sharding.h | 186 std::vector<int64> TileOffsetForDevice(const Shape& shape,
|
| D | hlo_sharding.cc | 263 std::vector<int64> HloSharding::TileOffsetForDevice(const Shape& shape, in TileOffsetForDevice() function in xla::HloSharding
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | xla_helpers.cc | 157 sharding->TileOffsetForDevice(*xla_shape, device); in RewriteLayoutWithShardedShape()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_compile_op_support.cc | 186 std::vector<int64> offset = sharding.TileOffsetForDevice(shape, device); in GetPerDeviceShape()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_util.cc | 191 sharding.TileOffsetForDevice(shape, partition_id); in MakeNonPaddedShapeForGivenPartition()
|