Home
last modified time | relevance | path

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

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/experimental/xla_sharding/
H A 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
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A 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
H A 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
H A 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
H A Dhlo_sharding.cc125 const Array<int64_t>& tile_assignment, in Subgroup()
694 Array<int64_t> tile_assignment( in FromProto() local
H A Dhlo_sharding.h289 const Array<int64_t>& tile_assignment() const { return tile_assignment_; } in tile_assignment() function
H A Dhlo_parser_test.cc3236 Array<int64_t> tile_assignment({2, 2, 2, 2}); local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/
H A Dsharding_builder.cc42 const TileAssignment& tile_assignment) { in Tile()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/
H A Dxla_compiler_test.cc1844 xla::Array<int64_t> tile_assignment({2}); in TEST_F() local