| 1053 | } |
| 1054 | |
| 1055 | Status AvgPoolShape(shape_inference::InferenceContext* c) { |
| 1056 | string data_format_str; |
| 1057 | TensorFormat data_format; |
| 1058 | Status s = c->GetAttr("data_format", &data_format_str); |
| 1059 | if (s.ok()) { |
| 1060 | FormatFromString(data_format_str, &data_format); |
| 1061 | } else { |
| 1062 | data_format = FORMAT_NHWC; |
| 1063 | } |
| 1064 | |
| 1065 | const int rank = (data_format == FORMAT_NCHW_VECT_C) ? 5 : 4; |
| 1066 | ShapeHandle input_shape; |
| 1067 | TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &input_shape)); |
| 1068 | |
| 1069 | TF_RETURN_IF_ERROR( |
| 1070 | CheckFormatConstraintsOnShape(data_format, input_shape, "input", c)); |
| 1071 | |
| 1072 | std::vector<int32> strides; |
| 1073 | TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides)); |
| 1074 | if (strides.size() != 4) { |
| 1075 | return errors::InvalidArgument( |
| 1076 | "AvgPool requires the stride attribute to contain 4 values, but got: ", |
| 1077 | strides.size()); |
| 1078 | } |
| 1079 | |
| 1080 | std::vector<int32> kernel_sizes; |
| 1081 | TF_RETURN_IF_ERROR(c->GetAttr("ksize", &kernel_sizes)); |
| 1082 | if (kernel_sizes.size() != 4) { |
| 1083 | return errors::InvalidArgument( |
| 1084 | "AvgPool requires the ksize attribute to contain 4 values, but got: ", |
| 1085 | kernel_sizes.size()); |
| 1086 | } |
| 1087 | |
| 1088 | int32 stride_rows = GetTensorDim(strides, data_format, 'H'); |
| 1089 | int32 stride_cols = GetTensorDim(strides, data_format, 'W'); |
| 1090 | int32 kernel_rows = GetTensorDim(kernel_sizes, data_format, 'H'); |
| 1091 | int32 kernel_cols = GetTensorDim(kernel_sizes, data_format, 'W'); |
| 1092 | |
| 1093 | constexpr int num_spatial_dims = 2; |
| 1094 | DimensionHandle batch_size_dim = c->Dim( |
| 1095 | input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'N')); |
| 1096 | DimensionHandle in_rows_dim = c->Dim( |
| 1097 | input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'H')); |
| 1098 | DimensionHandle in_cols_dim = c->Dim( |
| 1099 | input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'W')); |
| 1100 | DimensionHandle depth_dim = c->Dim( |
| 1101 | input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'C')); |
| 1102 | |
| 1103 | Padding padding; |
| 1104 | TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding)); |
| 1105 | |
| 1106 | // TODO(mrry,shlens): Raise an error if the stride would cause |
| 1107 | // information in the input to be ignored. This will require a change |
| 1108 | // in the kernel implementation. |
| 1109 | |
| 1110 | DimensionHandle output_rows, output_cols; |
| 1111 | TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims( |
| 1112 | c, in_rows_dim, kernel_rows, stride_rows, padding, &output_rows)); |
no test coverage detected