Home
last modified time | relevance | path

Searched defs:group_dim_shards (Results 1 – 3 of 3) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dgather_scatter_handler.cc213 absl::InlinedVector<int64_t, 1> group_dim_shards = { in PartitionGatherIndexPassthroughPartition() local
319 absl::InlinedVector<int64_t, 1> group_dim_shards = { in PartitionGatherPassthroughOperand() local
H A Ddot_handler.cc2329 std::vector<int64_t> group_dim_shards = { in GetNonContractingPartitionGroupedShardingForOtherOperand() local
2472 std::vector<int64_t> group_dim_shards = { in GetDotGroupPartitionContractingOutputShardings() local
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dhlo_sharding_util.cc1628 std::vector<int64_t> group_dim_shards(group_dims.size(), 1); in GroupShardingOnDims() local
1635 absl::Span<const int64_t> group_dim_shards, in GroupShardingOnDims()