| 202 | // before replicating this approach elsewhere. |
| 203 | #ifndef TF_LITE_STATIC_MEMORY |
| 204 | TfLiteStatus CalculateShapeForBroadcast(TfLiteContext* context, |
| 205 | const TfLiteTensor* input1, |
| 206 | const TfLiteTensor* input2, |
| 207 | TfLiteIntArray** output_shape) { |
| 208 | int64_t dims1 = NumDimensions(input1); |
| 209 | int64_t dims2 = NumDimensions(input2); |
| 210 | int64_t out_dims = std::max(dims1, dims2); |
| 211 | if (NumElements(input1) == 0) { |
| 212 | *output_shape = TfLiteIntArrayCopy(input1->dims); |
| 213 | return kTfLiteOk; |
| 214 | } |
| 215 | std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)> shape( |
| 216 | TfLiteIntArrayCreate(out_dims), TfLiteIntArrayFree); |
| 217 | for (int i = 0; i < out_dims; ++i) { |
| 218 | int64_t d1 = i >= dims1 ? 1 : SizeOfDimension(input1, dims1 - i - 1); |
| 219 | int64_t d2 = i >= dims2 ? 1 : SizeOfDimension(input2, dims2 - i - 1); |
| 220 | TF_LITE_ENSURE(context, d1 == d2 || d1 == 1 || d2 == 1); |
| 221 | shape->data[out_dims - i - 1] = std::max(d1, d2); |
| 222 | } |
| 223 | *output_shape = shape.release(); |
| 224 | return kTfLiteOk; |
| 225 | } |
| 226 | #endif // TF_LITE_STATIC_MEMORY |
| 227 | |
| 228 | } // namespace tflite |