Method Compile

tensorflow/compiler/tf2xla/kernels/pooling_ops.cc:169–188

  }

  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 pooling =
        xla::MaxPool(ctx->Input(0), ksize, stride, padding_,
                     XlaTensorFormat(data_format_, input_shape.dims() - 2));
    ctx->SetOutput(0, pooling);
  }
};

class MaxPool2DOp : public MaxPoolOp {