| 24 | }; |
| 25 | |
| 26 | std::shared_ptr<AvgPoolOpKernelCache> CreateAvgOpKernelCache(user_op::KernelCacheContext* ctx, |
| 27 | const int32_t& dim) { |
| 28 | const Shape& x_shape = ctx->TensorDesc4ArgNameAndIndex("x", 0)->shape(); |
| 29 | const std::string& data_format = ctx->Attr<std::string>("data_format"); |
| 30 | const std::vector<int32_t>& padding = ctx->Attr<std::vector<int32_t>>("padding"); |
| 31 | const std::vector<int32_t>& kernel_size = ctx->Attr<std::vector<int32_t>>("kernel_size"); |
| 32 | const std::vector<int32_t>& stride = ctx->Attr<std::vector<int32_t>>("stride"); |
| 33 | const bool ceil_mode = ctx->Attr<bool>("ceil_mode"); |
| 34 | const bool count_include_pad = ctx->Attr<bool>("count_include_pad"); |
| 35 | const int32_t divisor_override = ctx->Attr<int32_t>("divisor_override"); |
| 36 | |
| 37 | AvgPoolParams3D params_3d = |
| 38 | AvgPoolParams3D(dim, x_shape, data_format, padding, kernel_size, stride, ceil_mode, |
| 39 | count_include_pad, divisor_override); |
| 40 | std::shared_ptr<AvgPoolOpKernelCache> cache(new AvgPoolOpKernelCache(params_3d)); |
| 41 | return cache; |
| 42 | } |
| 43 | |
| 44 | template<typename T, typename IDX> |
| 45 | struct AvgPoolKernelUtil<DeviceType::kCPU, T, IDX> { |
no test coverage detected