void SetOutput(OpKernelContext* context, int ragged_rank,
                 const vector<INDEX_TYPE>& output_index,
                 Tensor* output_tensor) override {
    // Note: it's ok to use OP_REQUIRES_OK (rather than TF_RETURN_IF_ERROR)
    // in this function, but only because it's the last thing we do before
    // returning from Compute().

    if (output_tensor->NumElements() == 0) return;

    const auto& values_tensor = context->input(kValueInputIndex);
    const VALUE_TYPE* values_base = values_tensor.flat<VALUE_TYPE>().data();
    const auto& default_value_tensor = context->input(kDefaultValueInputIndex);
    VALUE_TYPE* output_base = output_tensor->flat<VALUE_TYPE>().data();

    TensorShape element_shape = output_tensor->shape();
    element_shape.RemoveDimRange(0, ragged_rank + 1);
    int value_element_size = element_shape.num_elements();
    size_t output_index_size = output_index.size();

    // Broadcast the default value to value_element_size.  (We can skip this
    // if default_value_tensor.NumElements() == 1, since we use std::fill
    // when that's true.)
    const VALUE_TYPE* default_value =
        default_value_tensor.flat<VALUE_TYPE>().data();
    Tensor bcast_default;  // Temporary tensor for result of broadcast
    if (default_value_tensor.NumElements() != value_element_size &&
        default_value_tensor.NumElements() != 1) {
      const auto& src_shape = default_value_tensor.shape();
      BCast bcast(BCast::FromShape(src_shape), BCast::FromShape(element_shape),
                  /*fewer_dims_optimization=*/true);
      // Note: bcast should always be valid, since we rejected any incompatible
      // shapes when we called ValidateDefaultValueShape().
      OP_REQUIRES(context, bcast.IsValid(),
                  errors::InvalidArgument("Error broadcasting default_value"));
      OP_REQUIRES_OK(context,
                     context->allocate_temp(default_value_tensor.dtype(),
                                            element_shape, &bcast_default));
      const CPUDevice& device = context->eigen_device<CPUDevice>();
      functor::BroadcastTo<CPUDevice, VALUE_TYPE>()(
          device, context, bcast_default, element_shape, default_value_tensor,
          src_shape, bcast);
      default_value = bcast_default.flat<VALUE_TYPE>().data();
    }

    // Loop through the output_index vector, finding contiguous regions that
    // should be copied.  Once we find the end of a contiguous region, copy it
    // and add any necessary padding (with default_value).
    INDEX_TYPE src_start = 0;  // Start of contiguous region (in values)
    INDEX_TYPE dst_start = 0;  // Destination for contiguous region (in output)
    INDEX_TYPE dst_end = 0;    // Destination for contiguous region (in output)
    for (int src_i = 0; src_i <= output_index_size; ++src_i) {
      // dst_i is the destination where the value at src_i should be copied.
      INDEX_TYPE dst_i = src_i < output_index_size ? output_index[src_i] : -1;

      // If we're still in a contiguous region, then update dst_end go to the
      // next src_i.
      if (dst_i == dst_end) {
        ++dst_end;
        continue;
      }

      // We found the end of contiguous region.  This can be because we found
      // a gap (dst_i > dst_end), or a source value that shouldn't be copied
      // because it's out-of-bounds (dst_i == -1), or the end of the tensor
      // (dst_i = -1).
      if (dst_start < dst_end) {
        // Copy the contiguous region.
        const VALUE_TYPE* src = values_base + src_start * value_element_size;
        VALUE_TYPE* dst = output_base + dst_start * value_element_size;
        INDEX_TYPE nvals = (dst_end - dst_start) * value_element_size;
        copy_array<VALUE_TYPE, INDEX_TYPE>(dst, src, nvals);
      }

      // Add any necessary padding (w/ default_value).
      if (src_i >= output_index_size) {
        // We reached the end of values: pad to the end of output.
        size_t output_size = output_tensor->NumElements();
        dst_i = output_size / value_element_size;
      }
      if (dst_i > dst_end) {
        if (default_value_tensor.NumElements() == 1) {
          std::fill(output_base + dst_end * value_element_size,
                    output_base + dst_i * value_element_size, *default_value);
          dst_end = dst_i;
        } else {
          while (dst_i > dst_end) {
            VALUE_TYPE* dst = output_base + dst_end * value_element_size;
            copy_array<VALUE_TYPE, INDEX_TYPE>(dst, default_value,
                                               value_element_size);
            ++dst_end;
          }
        }
      }

      // Update indices.
      if (dst_i < 0) {
        // src_i should be skipped -- leave it out of the contiguous region.
        src_start = src_i + 1;
        dst_start = dst_end;
      } else {
        // src_i should be copied -- include it in the contiguous region.
        src_start = src_i;
        dst_start = dst_end;
        dst_end = dst_start + 1;
      }
    }
  }