Searched refs:index_rank (Results 1 – 3 of 3) sorted by relevance
552 int64 index_rank = hlo->operand(1)->shape().rank(); in GatherIndexSharding() local555 if (index_rank > index_tile_assignment_dims.size()) { in GatherIndexSharding()736 const int64 index_rank = hlo.operand(1)->shape().rank(); in ScatterEffectiveIndexSharding() local737 std::vector<int64> slice_starts(index_rank, 0LL), slice_limits(index_rank); in ScatterEffectiveIndexSharding()738 for (int64 i = 0; i < index_rank; ++i) { in ScatterEffectiveIndexSharding()
1977 int64_t index_rank = index_ty.getRank(); in Verify() local1981 data_shape.take_front(index_rank)))) in Verify()1987 ArrayRef<int64_t> item_shape = data_shape.drop_front(index_rank); in Verify()
2255 int64 index_rank = c->Value(index_rank_dim); in ValidateSparseTensor() local2257 if (index_rank != shape_rank) { in ValidateSparseTensor()2258 return errors::InvalidArgument("Index rank (", index_rank, in ValidateSparseTensor()