Commit bf7dcf4d authored by Alina Sbirlea's avatar Alina Sbirlea Committed by TensorFlower Gardener
Browse files

Correct wrapping code generation for DynamicUpdateSlice.

Currently if slice_index falls outside of the input, those values are not updated.
With wrapping behavior:
for each rank:
  // Index for update:
  index
  // Indexes for input:
  slice_start_index = slice_start_index % input_size //in-bounds
  slice_limit_index = slice_start_index + update_size //may be out of bounds
  // If in bounds:
  if (slice_limit_index <= input_size) {
    slice intersects = (index > slice_start_index) && (index <= slice_limit_index)
  }
  // If out of bounds:
  else {
    slice intersects = (index > slice_start_index) || (index <= slice_limit_index%input_size)
  }

if (slice intersects)
  for all ranks, compute index for intersection:
     if slice_index was in bounds
       intersection index [rank] = index - slice_start_index
     else (out of bounds) && (index <= slice_limit_index%input_size)
       intersection index [rank] = index + update_size - slice_start_index
  return update [intersection index]
else
  return input [index]

The above algorithms does not account for negative indexes given as inputs.

PiperOrigin-RevId: 174915327
parent 6b74eb0b
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment