| 432 | |
| 433 | template <typename Device, typename T> |
| 434 | Status ReduceOperand(OpKernelContext* ctx, const Tensor& input, |
| 435 | const std::vector<DimensionType>& label_types, |
| 436 | const LabelCounts& label_counts, Labels* labels, |
| 437 | Labels* free_labels, bool* swap_free_and_contract, |
| 438 | Tensor* output) { |
| 439 | // Find the permutation to transpose the input dimensions in the order of |
| 440 | // DimensionType; i.e. batch, free, contract and reduce dimensions. This |
| 441 | // makes it more convenient to invoke Reduce/Contract operations. |
| 442 | std::vector<int> permutation(input.dims()); |
| 443 | absl::c_iota(permutation, 0); |
| 444 | Tensor input_transposed; |
| 445 | // Check if we can avoid the transpose. We need to flip the adj_x (or adj_y) |
| 446 | // flag during BatchMatMul. This is an extra optimization not necessary for |
| 447 | // correctness. |
| 448 | if (ShouldSwapFreeAndContract(*labels, label_types)) { |
| 449 | *swap_free_and_contract = true; |
| 450 | } else { |
| 451 | absl::c_sort(permutation, [&](int i, int j) { |
| 452 | int label_i = (*labels)[i]; |
| 453 | int label_j = (*labels)[j]; |
| 454 | return std::tie(label_types[label_i], label_i) < |
| 455 | std::tie(label_types[label_j], label_j); |
| 456 | }); |
| 457 | } |
| 458 | // Transpose the input so that DimensionTypes are in order. |
| 459 | TF_RETURN_IF_ERROR( |
| 460 | TransposeOperand<Device, T>(ctx, input, permutation, &input_transposed)); |
| 461 | PermuteLabels(permutation, labels); |
| 462 | |
| 463 | // Take the generalized diagonal for dimensions with repeated axis labels. |
| 464 | Tensor input_deduped; |
| 465 | labels->erase(std::unique(labels->begin(), labels->end()), labels->end()); |
| 466 | TF_RETURN_IF_ERROR( |
| 467 | StrideOrInflate<Device, T>(ctx, input_transposed, *labels, label_counts, |
| 468 | false /* should_inflate */, &input_deduped)); |
| 469 | |
| 470 | // Reshape denotes the rank-5 shape [broadcast, batch, free, contract, reduce] |
| 471 | // where we've compacted the dimensions of each DimensionType. |
| 472 | gtl::InlinedVector<int64, 5> reshape(5, 1); |
| 473 | // The output shape is [batch shape] + [free size, contract size] |
| 474 | // That is, the batch shape is preserved (for broadcasting while contracting) |
| 475 | // while the free dims and contract dims are compressed to one dimension each. |
| 476 | TensorShape output_shape; |
| 477 | for (int label_idx = 0; label_idx < labels->size(); ++label_idx) { |
| 478 | const int label = labels->at(label_idx); |
| 479 | int64 dim = input_deduped.dim_size(label_idx); |
| 480 | if (label_types[label] == kBroadcasting || label_types[label] == kBatch) { |
| 481 | output_shape.AddDim(dim); |
| 482 | } else if (label_types[label] == kFree) { |
| 483 | free_labels->push_back(label); |
| 484 | } |
| 485 | reshape[label_types[label]] *= dim; |
| 486 | } |
| 487 | if (*swap_free_and_contract) std::swap(reshape[kFree], reshape[kContract]); |
| 488 | output_shape.AddDim(reshape[kFree]); |
| 489 | output_shape.AddDim(reshape[kContract]); |
| 490 | |
| 491 | if (reshape[kReduce] == 1) { // No need to actually reduce. |
nothing calls this directly
no test coverage detected