op, input.getType(), scalar_zero,
        ArrayRef<Value>({slice1, expanded_item, slice2}));
    return success();
  }

  // This function rewrites the original op into a XLA DynamicUpdateSlice op.
  // |item| is expanded to have the same dimension as input_handle and
  // |index| is expanded to [index, 0, 0, ...] as the indices to input_handle.
  // On a high level, it's doing something like: