Method Compile

tensorflow/compiler/tf2xla/kernels/pooling_ops.cc:221–249

  }

  void Compile(XlaOpKernelContext* ctx) override {
    auto ksize_or_error = GetKernelSize(ctx);
    OP_REQUIRES_OK(ctx, ksize_or_error.status());
    std::vector<int64> ksize = ksize_or_error.ValueOrDie();

    auto stride_or_error = GetStride(ctx);
    OP_REQUIRES_OK(ctx, stride_or_error.status());
    std::vector<int64> stride = stride_or_error.ValueOrDie();

    const TensorShape input_shape = ctx->InputShape(0);
    OP_REQUIRES(ctx, input_shape.dims() == num_dims(),
                errors::InvalidArgument("Input to ", type_string(),
                                        " operator must have ", num_dims(),
                                        " dimensions"));

    auto xla_data_format =
        XlaTensorFormat(data_format_, input_shape.dims() - 2);
    auto spatial_padding = MakeSpatialPadding(
        input_shape.dim_sizes(), ksize, stride, padding_, xla_data_format);

    // Convert the input to the reduction type.
    auto converted_input =
        ConvertElementType(ctx->Input(0), xla_reduction_type_);
    auto pooling =
        xla::AvgPool(converted_input, ksize, stride, spatial_padding,
                     xla_data_format, padding_ == xla::Padding::kValid);
    // Convert the pooling result back to the input type before returning it.
    ctx->SetOutput(0, ConvertElementType(pooling, ctx->input_xla_type(0)));
  }
};

class AvgPool2DOp : public AvgPoolOp {