Method UpdateSlice

tensorflow/compiler/xla/array.h:471–488

Performs the equivalent of a DynamicUpdateSlice in-place on this array.

  // Performs the equivalent of a DynamicUpdateSlice in-place on this array.
  void UpdateSlice(const Array<T>& from,
                   absl::Span<const int64> start_indices) {
    CHECK_EQ(from.num_dimensions(), num_dimensions());
    std::vector<int64> limit_indices;
    std::transform(start_indices.begin(), start_indices.end(),
                   from.dimensions().begin(), std::back_inserter(limit_indices),
                   std::plus<int64>{});
    std::vector<int64> index(sizes_.size());
    int64 from_i = 0;
    for (int64 i = 0; i < num_elements(); ++i, next_index(&index)) {
      if (array_impl::all_inside_range(index, start_indices, limit_indices)) {
        // Even though the bounds of from are different to our bounds, we're
        // iterating in the same order. So we can simply write successive linear
        // indices instead of recalculating a multi-dimensional index.
        values_[i] = from.values_[from_i++];
      }
    }
  }

  // Performs an in-place reshape, modifying the dimensions but not the
  // underlying data.