Searched refs:updates_rank (Results 1 – 5 of 5) sorted by relevance
143 int64 updates_rank = updates_shape.rank(); in XlaScatter() local156 if (updates_rank == 0 && expected_updates_rank != 0) { in XlaScatter()159 updates_rank = updates_shape.rank(); in XlaScatter()162 if (updates_rank > 0) { in XlaScatter()163 for (int64 i = (updates_rank - num_window_dims_in_updates); in XlaScatter()164 i < updates_rank; ++i) { in XlaScatter()
97 const int64 updates_rank = updates->shape().rank(); in PermuteScatterAndWindowDims() local98 permutation.reserve(updates_rank); in PermuteScatterAndWindowDims()100 for (int64 i = 0; i < updates_rank; ++i) { in PermuteScatterAndWindowDims()
2083 int64 updates_rank = updates_shape.dimensions_size();2084 std::vector<int64> index_base(updates_rank, 0);2085 std::vector<int64> index_count(updates_rank, 1);2086 for (int64 i = 0; i < updates_rank; i++) {2094 std::vector<int64>(updates_rank, 1)};2102 int64 updates_rank = updates_shape.dimensions_size();2103 std::vector<int64> index_base(updates_rank, 0);2104 std::vector<int64> index_count(updates_rank, 1);2105 for (int64 i = 0; i < updates_rank; i++) {2113 std::vector<int64>(updates_rank, 1)};
3459 const int64 updates_rank = updates_shape.rank(); in ValidateScatterDimensionNumbers() local3461 if (window_dim < 0 || window_dim >= updates_rank) { in ValidateScatterDimensionNumbers()3465 updates_rank, window_dim); in ValidateScatterDimensionNumbers()
3936 int64_t updates_rank = updates_ty.getRank(); in matchAndRewrite() local3940 GetI64ElementsAttrForSeq(updates_rank - window_dims, updates_rank, in matchAndRewrite()