static */
| 1094 | } |
| 1095 | |
| 1096 | /* static */ StatusOr<Shape> ShapeInference::InferMapShape( |
| 1097 | absl::Span<const Shape* const> arg_shapes, const ProgramShape& to_apply, |
| 1098 | absl::Span<const int64> dimensions) { |
| 1099 | if (arg_shapes.empty()) { |
| 1100 | return InvalidArgument("Map expects at least one argument."); |
| 1101 | } |
| 1102 | |
| 1103 | // All arguments must have the same shape. |
| 1104 | const Shape* arg_shape = arg_shapes[0]; |
| 1105 | for (size_t i = 1; i < arg_shapes.size(); ++i) { |
| 1106 | TF_RETURN_IF_ERROR(ExpectArray(*arg_shapes[i], "operand of map")); |
| 1107 | |
| 1108 | if (ShapeUtil::CompatibleIgnoringFpPrecision(*arg_shapes[i], *arg_shape)) { |
| 1109 | continue; |
| 1110 | } |
| 1111 | if (ShapeUtil::SameElementTypeIgnoringFpPrecision(*arg_shapes[i], |
| 1112 | *arg_shape)) { |
| 1113 | if (ShapeUtil::IsScalar(*arg_shapes[i])) { |
| 1114 | continue; |
| 1115 | } |
| 1116 | if (ShapeUtil::IsScalar(*arg_shape)) { |
| 1117 | arg_shape = arg_shapes[i]; |
| 1118 | continue; |
| 1119 | } |
| 1120 | } |
| 1121 | |
| 1122 | std::vector<string> pieces; |
| 1123 | for (const Shape* shape : arg_shapes) { |
| 1124 | pieces.push_back(ShapeUtil::HumanString(*shape)); |
| 1125 | } |
| 1126 | return InvalidArgument( |
| 1127 | "Map operation requires all operands to have the same shape; got: " |
| 1128 | "%s.", |
| 1129 | StrJoin(pieces, ", ")); |
| 1130 | } |
| 1131 | |
| 1132 | // Check that dimensions.size == arg_shape.dimensions_size() (we currently |
| 1133 | // only support mapping across all dimensions: i.e. scalar map functions). |
| 1134 | if (dimensions.size() != arg_shape->dimensions_size()) { |
| 1135 | return InvalidArgument( |
| 1136 | "Map applied to a subset of dimensions currently not supported: " |
| 1137 | "arg_dimension_size: %d, requested_map_dimensions_size: %u.", |
| 1138 | arg_shape->dimensions_size(), dimensions.size()); |
| 1139 | } |
| 1140 | |
| 1141 | // Check that requested map dimensions numbers are monotonically increasing. |
| 1142 | for (int i = 0; i < dimensions.size(); ++i) { |
| 1143 | if (dimensions[i] != i) { |
| 1144 | return InvalidArgument( |
| 1145 | "Map requires monotonically increasing dimension numbers; got: %s.", |
| 1146 | StrJoin(dimensions, ", ")); |
| 1147 | } |
| 1148 | } |
| 1149 | |
| 1150 | // The applied function's arity equals the number of arguments. |
| 1151 | if (arg_shapes.size() != to_apply.parameters_size()) { |
| 1152 | return InvalidArgument( |
| 1153 | "Map applied function arity must match number of arguments; got: " |
nothing calls this directly
no test coverage detected