Searched refs:updates_rank (Results 1 – 4 of 4) 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()
2060 int64 updates_rank = updates_shape.dimensions_size(); in IterationSpaceForUpdateScatterIndices() local2061 std::vector<int64> index_base(updates_rank, 0); in IterationSpaceForUpdateScatterIndices()2062 std::vector<int64> index_count(updates_rank, 1); in IterationSpaceForUpdateScatterIndices()2063 for (int64 i = 0; i < updates_rank; i++) { in IterationSpaceForUpdateScatterIndices()2071 std::vector<int64>(updates_rank, 1)}; in IterationSpaceForUpdateScatterIndices()2079 int64 updates_rank = updates_shape.dimensions_size(); in IterationSpaceForUpdateWindowIndices() local2080 std::vector<int64> index_base(updates_rank, 0); in IterationSpaceForUpdateWindowIndices()2081 std::vector<int64> index_count(updates_rank, 1); in IterationSpaceForUpdateWindowIndices()2082 for (int64 i = 0; i < updates_rank; i++) { in IterationSpaceForUpdateWindowIndices()2090 std::vector<int64>(updates_rank, 1)}; in IterationSpaceForUpdateWindowIndices()
3084 const int64 updates_rank = updates_shape.rank(); in ValidateScatterDimensionNumbers() local3086 if (window_dim < 0 || window_dim >= updates_rank) { in ValidateScatterDimensionNumbers()3090 updates_rank, window_dim); in ValidateScatterDimensionNumbers()