void Compute(OpKernelContext* context) override {
    const Tensor& filter = context->input(1);
    const TensorShape& filter_shape = filter.shape();

    const Tensor& out_backprop = context->input(2);
    const TensorShape& out_backprop_shape = out_backprop.shape();

    TensorShape input_shape;
    if (takes_shape_) {
      const Tensor& input_sizes = context->input(0);
      // tensor::MakeShape is able to handle both DT_INT32 and DT_INT64 for
      // input_sizes.
      OP_REQUIRES_OK(context, tensor::MakeShape(input_sizes, &input_shape));
    } else {
      input_shape = context->input(0).shape();
    }

    OP_REQUIRES(context, input_shape.dims() == 5,
                errors::InvalidArgument("input tensor must have 5 dimensions"));
    OP_REQUIRES(
        context, filter_shape.dims() == 5,
        errors::InvalidArgument("filter_sizes tensor must have 5 dimensions"));
    OP_REQUIRES(
        context, out_backprop_shape.dims() == 5,
        errors::InvalidArgument("out_backprop tensor must have 5 dimensions"));
    OP_REQUIRES(
        context, input_shape.dim_size(4) == filter_shape.dim_size(3),
        errors::InvalidArgument("input and filter_sizes must have the same "
                                "number of channels. Got ",
                                input_shape.dim_size(4), " for input and ",
                                filter_shape.dim_size(3), " for filter_sizes"));
    OP_REQUIRES(
        context, out_backprop_shape.dim_size(4) == filter_shape.dim_size(4),
        errors::InvalidArgument("out_backprop and filter_sizes must have the "
                                "same number of channels. Got ",
                                out_backprop_shape.dim_size(4),
                                " for out_backprop and ",
                                filter_shape.dim_size(4), " for filter_sizes"));

    ConvBackpropDimensions dims;
    OP_REQUIRES_OK(context, ConvBackpropComputeDimensions(
                                "Conv3DBackpropInputOp", /*num_spatial_dims=*/3,
                                input_shape, filter_shape, out_backprop_shape,
                                stride_, padding_, data_format_, &dims));

    Tensor* in_backprop;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_shape, &in_backprop));

    int64 top_pad_planes, bottom_pad_planes;
    int64 top_pad_rows, bottom_pad_rows;
    int64 left_pad_cols, right_pad_cols;

    OP_REQUIRES_OK(context, GetWindowedOutputSizeVerbose(
                                dims.spatial_dims[0].input_size,
                                dims.spatial_dims[0].filter_size,
                                dims.spatial_dims[0].stride, padding_,
                                &dims.spatial_dims[0].output_size,
                                &top_pad_planes, &bottom_pad_planes));
    OP_REQUIRES_OK(context, GetWindowedOutputSizeVerbose(
                                dims.spatial_dims[1].input_size,
                                dims.spatial_dims[1].filter_size,
                                dims.spatial_dims[1].stride, padding_,
                                &dims.spatial_dims[1].output_size,
                                &top_pad_rows, &bottom_pad_rows));
    OP_REQUIRES_OK(context, GetWindowedOutputSizeVerbose(
                                dims.spatial_dims[2].input_size,
                                dims.spatial_dims[2].filter_size,
                                dims.spatial_dims[2].stride, padding_,
                                &dims.spatial_dims[2].output_size,
                                &left_pad_cols, &right_pad_cols));

    // TODO(ezhulenev): Extract work size and shard estimation to shared
    // functions in conv_grad_ops, and update 2d convolution backprop.

    // The total dimension size of each kernel.
    const int64 filter_total_size =
        dims.spatial_dims[0].filter_size * dims.spatial_dims[1].filter_size *
        dims.spatial_dims[2].filter_size * dims.in_depth;

    // The output image size is the spatial size of the output.
    const int64 output_image_size = dims.spatial_dims[0].output_size *
                                    dims.spatial_dims[1].output_size *
                                    dims.spatial_dims[2].output_size;

    const auto cache_sizes = Eigen::internal::CacheSizes();
    const ptrdiff_t l3_cache_size = cache_sizes.m_l3;

    // Use L3 cache size as target working set size.
    const size_t target_working_set_size = l3_cache_size / sizeof(T);

    // Calculate size of matrices involved in MatMul: C = A x B.
    const int64 size_A = output_image_size * dims.out_depth;

    const int64 size_B = filter_total_size * dims.out_depth;

    const int64 size_C = output_image_size * filter_total_size;

    const int64 work_unit_size = size_A + size_B + size_C;

    auto worker_threads = *(context->device()->tensorflow_cpu_worker_threads());

    // Use parallel tensor contractions if there is no batching.
    //
    // Compared to Conv2D code, this version is missing work size estimation. In
    // benchmarks I didn't find a case when it's beneficial to run parallel
    // contraction compared to sharding and matmuls.
    const bool use_parallel_contraction = dims.batch_size == 1;

    OP_REQUIRES(
        context, work_unit_size > 0,
        errors::InvalidArgument("input, filter_sizes and out_backprop tensors "
                                "must all have at least 1 element"));

    const size_t shard_size =
        use_parallel_contraction
            ? 1
            : (target_working_set_size + work_unit_size - 1) / work_unit_size;

    // Total number of elements in all the tensors used by this kernel.
    int64 total_tensor_elements = input_shape.num_elements() +
                                  filter_shape.num_elements() +
                                  out_backprop_shape.num_elements();

    // Shape of the temporary workspace buffer.
    TensorShape col_buffer_shape = {static_cast<int64>(shard_size),
                                    static_cast<int64>(output_image_size),
                                    static_cast<int64>(filter_total_size)};
    int64 col_buffer_elements = col_buffer_shape.num_elements();

    // If the temporary allocation overhead is too large, fallback on Eigen
    // implementation which requires much less memory.
    int64 col_buffer_overhead = col_buffer_elements / total_tensor_elements;
    if (col_buffer_overhead > kMaxTempAllocationOverhead) {
      VLOG(2) << "Fallback on Eigen implementation of Conv3DBackpropInputOp: "
                 "col_buffer_overhead="
              << col_buffer_overhead;

      functor::CuboidConvolutionBackwardInput<Device, T>()(
          context->eigen_device<Device>(),
          in_backprop->tensor<T, 5>(),                     // input_backward
          filter.tensor<T, 5>(),                           // filter
          out_backprop.tensor<T, 5>(),                     // output_backward
          static_cast<int>(dims.spatial_dims[0].stride),   // stride_planes
          static_cast<int>(dims.spatial_dims[1].stride),   // stride_rows
          static_cast<int>(dims.spatial_dims[2].stride));  // stride_cols

      return;
    }

    Tensor col_buffer;
    OP_REQUIRES_OK(context,
                   context->allocate_temp(DataTypeToEnum<T>::value,
                                          col_buffer_shape, &col_buffer));

    // The input offset corresponding to a single input image.
    const int64 input_offset = dims.spatial_dims[0].input_size *
                               dims.spatial_dims[1].input_size *
                               dims.spatial_dims[2].input_size * dims.in_depth;

    // The output offset corresponding to a single output image.
    const int64 output_offset =
        dims.spatial_dims[0].output_size * dims.spatial_dims[1].output_size *
        dims.spatial_dims[2].output_size * dims.out_depth;

    const T* filter_data = filter.template flat<T>().data();
    T* col_buffer_data = col_buffer.template flat<T>().data();
    const T* out_backprop_data = out_backprop.template flat<T>().data();

    auto in_backprop_flat = in_backprop->template flat<T>();
    T* input_backprop_data = in_backprop_flat.data();
    in_backprop_flat.device(context->eigen_device<Device>()) =
        in_backprop_flat.constant(T(0));

    if (use_parallel_contraction) {
      typedef Eigen::TensorMap<Eigen::Tensor<T, 2, Eigen::RowMajor>,
                               Eigen::Unaligned>
          TensorMap;
      typedef Eigen::TensorMap<Eigen::Tensor<const T, 2, Eigen::RowMajor>,
                               Eigen::Unaligned>
          ConstTensorMap;

      // Initialize contraction dims (we need to transpose 'B' below).
      Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1> contract_dims;
      contract_dims[0].first = 1;
      contract_dims[0].second = 1;

      for (int image_id = 0; image_id < dims.batch_size; ++image_id) {
        // Compute gradient into col_buffer.
        TensorMap C(col_buffer_data, output_image_size, filter_total_size);

        ConstTensorMap A(out_backprop_data + output_offset * image_id,
                         output_image_size, dims.out_depth);
        ConstTensorMap B(filter_data, filter_total_size, dims.out_depth);

        C.device(context->eigen_cpu_device()) = A.contract(B, contract_dims);

        Col2im<T>(col_buffer_data, dims.in_depth,
                  // Input spatial dimensions.
                  dims.spatial_dims[0].input_size,  // input planes
                  dims.spatial_dims[1].input_size,  // input rows
                  dims.spatial_dims[2].input_size,  // input cols
                  // Filter spatial dimensions.
                  dims.spatial_dims[0].filter_size,  // filter planes
                  dims.spatial_dims[1].filter_size,  // filter rows
                  dims.spatial_dims[2].filter_size,  // filter cols
                  // Spatial padding.
                  top_pad_planes, top_pad_rows, left_pad_cols,
                  bottom_pad_planes, bottom_pad_rows, right_pad_cols,
                  // Spatial striding.
                  dims.spatial_dims[0].stride,  // stride planes
                  dims.spatial_dims[1].stride,  // stride rows
                  dims.spatial_dims[2].stride,  // stride cols
                  input_backprop_data);

        input_backprop_data += input_offset;
      }
    } else {
      typedef Eigen::Map<
          Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>
          MatrixMap;
      typedef Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic,
                                             Eigen::RowMajor>>
          ConstMatrixMap;

      for (int image_id = 0; image_id < dims.batch_size;
           image_id += shard_size) {
        const int shard_limit =
            std::min(static_cast<int>(shard_size),
                     static_cast<int>(dims.batch_size) - image_id);

        auto shard = [&dims, &top_pad_planes, &top_pad_rows, &left_pad_cols,
                      &bottom_pad_planes, &bottom_pad_rows, &right_pad_cols,
                      &output_image_size, &filter_total_size,
                      &input_backprop_data, &col_buffer_data,
                      &out_backprop_data, &filter_data, &input_offset,
                      &output_offset, &size_C](int64 start, int64 limit) {
          for (int shard_id = start; shard_id < limit; ++shard_id) {
            T* im2col_buf = col_buffer_data + shard_id * size_C;
            T* input_data = input_backprop_data + shard_id * input_offset;
            const T* out_data = out_backprop_data + shard_id * output_offset;

            // Compute gradient into 'im2col_buf'.
            MatrixMap C(im2col_buf, output_image_size, filter_total_size);

            ConstMatrixMap A(out_data, output_image_size, dims.out_depth);
            ConstMatrixMap B(filter_data, filter_total_size, dims.out_depth);

            C.noalias() = A * B.transpose();

            Col2im<T>(im2col_buf, dims.in_depth,
                      // Input spatial dimensions.
                      dims.spatial_dims[0].input_size,  // input planes
                      dims.spatial_dims[1].input_size,  // input rows
                      dims.spatial_dims[2].input_size,  // input cols
                      // Filter spatial dimensions.
                      dims.spatial_dims[0].filter_size,  // filter planes
                      dims.spatial_dims[1].filter_size,  // filter rows
                      dims.spatial_dims[2].filter_size,  // filter cols
                      // Spatial padding.
                      top_pad_planes, top_pad_rows, left_pad_cols,
                      bottom_pad_planes, bottom_pad_rows, right_pad_cols,
                      // Spatial striding.
                      dims.spatial_dims[0].stride,  // stride planes
                      dims.spatial_dims[1].stride,  // stride rows
                      dims.spatial_dims[2].stride,  // stride cols
                      input_data);
          }
        };
        Shard(worker_threads.num_threads, worker_threads.workers, shard_limit,
              work_unit_size, shard);

        input_backprop_data += input_offset * shard_limit;
        out_backprop_data += output_offset * shard_limit;
      }
    }
  }