Method Compile

tensorflow/compiler/tf2xla/kernels/pooling_ops.cc:423–464

  int num_dims() const { return num_spatial_dims_ + 2; }

  void Compile(XlaOpKernelContext* ctx) override {
    TensorShape gradients_shape;
    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsShape(0, &gradients_shape));

    const TensorShape out_backprop_shape = ctx->InputShape(1);

    // For avgpooling, tensor_in_shape should have num_dims() dimensions.
    OP_REQUIRES(ctx, gradients_shape.dims() == num_dims(),
                errors::InvalidArgument("orig_input_shape must be ", num_dims(),
                                        "-dimensional"));

    // For avgpooling, out_backprop should have num_dims() dimensions.
    OP_REQUIRES(ctx, out_backprop_shape.dims() == num_dims(),
                errors::InvalidArgument("out_backprop must be ", num_dims(),
                                        "-dimensional"));

    auto out_backprop = ctx->Input(1);
    std::vector<int64> stride_int64s(stride_.begin(), stride_.end());
    xla::Padding xla_padding =
        (padding_ == VALID) ? xla::Padding::kValid : xla::Padding::kSame;
    xla::PrimitiveType xla_reduction_type;
    auto reduction_type = XlaHelpers::SumAccumulationType(ctx->input_type(1));
    OP_REQUIRES_OK(
        ctx, DataTypeToPrimitiveType(reduction_type, &xla_reduction_type));
    auto converted_out_backprop =
        xla::ConvertElementType(out_backprop, xla_reduction_type);
    auto xla_data_format =
        XlaTensorFormat(data_format_, gradients_shape.dims() - 2);
    auto padding_values =
        MakeSpatialPadding(gradients_shape.dim_sizes(), ksize_, stride_int64s,
                           xla_padding, xla_data_format);
    auto in_backprop =
        xla::AvgPoolGrad(converted_out_backprop, gradients_shape.dim_sizes(),
                         ksize_, stride_int64s, padding_values, xla_data_format,
                         /*counts_include_padding=*/padding_ == VALID);
    // Convert the pooling result back to the input type before returning it.
    xla::PrimitiveType xla_out_backprop_type;
    OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(ctx->input_type(1),
                                                &xla_out_backprop_type));
    ctx->SetOutput(0,
                   xla::ConvertElementType(in_backprop, xla_out_backprop_type));
  }

 protected:
  const int num_spatial_dims_;