TODO(mjanusz): Unify all conv/pooling shape functions.
| 865 | |
| 866 | // TODO(mjanusz): Unify all conv/pooling shape functions. |
| 867 | Status Conv3DShape(shape_inference::InferenceContext* c) { |
| 868 | ShapeHandle input_shape; |
| 869 | TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 5, &input_shape)); |
| 870 | ShapeHandle filter_shape; |
| 871 | TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 5, &filter_shape)); |
| 872 | |
| 873 | string data_format; |
| 874 | Status s = c->GetAttr("data_format", &data_format); |
| 875 | |
| 876 | std::vector<int32> dilations; |
| 877 | TF_RETURN_IF_ERROR(c->GetAttr("dilations", &dilations)); |
| 878 | |
| 879 | if (dilations.size() != 5) { |
| 880 | return errors::InvalidArgument( |
| 881 | "Conv3D requires the dilation attribute to contain 5 values, but got: ", |
| 882 | dilations.size()); |
| 883 | } |
| 884 | |
| 885 | std::vector<int32> strides; |
| 886 | TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides)); |
| 887 | if (strides.size() != 5) { |
| 888 | return errors::InvalidArgument( |
| 889 | "Conv3D requires the stride attribute to contain 5 values, but got: ", |
| 890 | strides.size()); |
| 891 | } |
| 892 | |
| 893 | int32 stride_planes, stride_rows, stride_cols; |
| 894 | int32 dilation_planes, dilation_rows, dilation_cols; |
| 895 | if (s.ok() && data_format == "NCDHW") { |
| 896 | // Convert input_shape to NDHWC. |
| 897 | auto dim = [&](char dimension) { |
| 898 | return c->Dim(input_shape, GetTensorDimIndex<3>(FORMAT_NCHW, dimension)); |
| 899 | }; |
| 900 | input_shape = |
| 901 | c->MakeShape({{dim('N'), dim('0'), dim('1'), dim('2'), dim('C')}}); |
| 902 | stride_planes = strides[2]; |
| 903 | stride_rows = strides[3]; |
| 904 | stride_cols = strides[4]; |
| 905 | dilation_planes = dilations[2]; |
| 906 | dilation_cols = dilations[3]; |
| 907 | dilation_rows = dilations[4]; |
| 908 | } else { |
| 909 | stride_planes = strides[1]; |
| 910 | stride_rows = strides[2]; |
| 911 | stride_cols = strides[3]; |
| 912 | dilation_planes = dilations[1]; |
| 913 | dilation_cols = dilations[2]; |
| 914 | dilation_rows = dilations[3]; |
| 915 | } |
| 916 | |
| 917 | DimensionHandle batch_size_dim = c->Dim(input_shape, 0); |
| 918 | DimensionHandle in_planes_dim = c->Dim(input_shape, 1); |
| 919 | DimensionHandle in_rows_dim = c->Dim(input_shape, 2); |
| 920 | DimensionHandle in_cols_dim = c->Dim(input_shape, 3); |
| 921 | DimensionHandle input_depth_dim = c->Dim(input_shape, 4); |
| 922 | |
| 923 | DimensionHandle filter_planes_dim = c->Dim(filter_shape, 0); |
| 924 | DimensionHandle filter_rows_dim = c->Dim(filter_shape, 1); |
nothing calls this directly
no test coverage detected