Home
last modified time | relevance | path

Searched defs:tile_assignment (Results 1 – 9 of 9) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/
Dxla_sharding.py71 def tile(cls, tile_assignment): argument
99 def subgroup_tile(cls, tile_assignment, subgroup_modes): argument
144 def partial_tile(cls, tile_assignment): argument
327 tile_assignment, argument
372 tile_assignment, argument
/external/tensorflow/tensorflow/compiler/xla/client/
Dsharding_builder.cc42 const TileAssignment& tile_assignment) { in Tile()
/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_sharding_util.cc339 Array<int64_t> tile_assignment = sharding.tile_assignment(); in TransposeSharding() local
567 Array<int64_t> tile_assignment(tile_dims); in ReshapeToTileDimension() local
725 Array<int64_t> tile_assignment = in GatherEffectiveOutputSharding() local
854 Array<int64_t> tile_assignment = in ScatterEffectiveIndexSharding() local
897 Array<int64_t> tile_assignment = in ScatterEffectiveDataSharding() local
944 Array<int64_t> tile_assignment = operand_sharding.tile_assignment(); in PassthroughOperandToGatherOutputOrScatterUpdate() local
999 Array<int64_t> tile_assignment = relevant_sharding.tile_assignment(); in PassthroughGatherOutputOrScatterUpdateToOperand() local
1049 Array<int64_t> tile_assignment = relevant_output_sharding.tile_assignment(); in GatherParallelDataOperandSharding() local
Dhlo_sharding_util_test.cc234 Array<int64_t> tile_assignment({1, 2, 2}); in TEST() local
254 Array<int64_t> tile_assignment({1, 2, 2, 2}); in TEST() local
275 Array<int64_t> tile_assignment({1, 2, 2, 2}); in TEST() local
Dhlo_sharding.h289 const Array<int64_t>& tile_assignment() const { return tile_assignment_; } in tile_assignment() function
Dsharding_propagation.cc599 const auto& tile_assignment = lhs->sharding().tile_assignment(); in InferConvolutionShardingFromOperands() local
1341 const Array<int64_t>& tile_assignment = user.sharding().tile_assignment(); in GetShardingFromUser() local
1409 const auto& tile_assignment = user.sharding().tile_assignment(); in GetShardingFromUser() local
1539 auto tile_assignment = user_sharding.tile_assignment(); in GetShardingFromUser() local
1652 const auto& tile_assignment = operand->sharding().tile_assignment(); in InferDynamicSliceOrDynamicUpdateSliceShardingFromOperands() local
1966 const auto& tile_assignment = in InferShardingFromOperands() local
Dhlo_sharding.cc125 const Array<int64_t>& tile_assignment, in Subgroup()
694 Array<int64_t> tile_assignment( in FromProto() local
Dhlo_parser_test.cc3236 Array<int64_t> tile_assignment({2, 2, 2, 2}); local
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_compiler_test.cc1844 xla::Array<int64_t> tile_assignment({2}); in TEST_F() local