| 573 | } |
| 574 | |
| 575 | void Compute(OpKernelContext* ctx) override { |
| 576 | OpInputList inputs; |
| 577 | OP_REQUIRES_OK(ctx, ctx->input_list("inputs", &inputs)); |
| 578 | |
| 579 | OperandLabels input_labels(input_labels_); |
| 580 | Labels output_labels(output_labels_); |
| 581 | std::vector<DimensionType> label_types(label_types_); |
| 582 | OperandLabelCounts input_label_counts(input_label_counts_); |
| 583 | LabelCounts output_label_counts(output_label_counts_); |
| 584 | LabelToDimSizes label_to_dim_sizes; |
| 585 | |
| 586 | OP_REQUIRES_OK(ctx, ProcessDimensions( |
| 587 | inputs, input_has_ellipsis_, output_has_ellipsis_, |
| 588 | &input_labels, &output_labels, &label_types, |
| 589 | &input_label_counts, &output_label_counts, |
| 590 | &label_to_dim_sizes)); |
| 591 | |
| 592 | // The reduction phase (a) sums across reduction dimensions, (b) takes |
| 593 | // generalized diagonals, and (c) reshapes it into shape |
| 594 | // [(broadcasting) batch shape] + [F,C] |
| 595 | // where F and C denote the total (compacted) size of free and contract |
| 596 | // dimensions, respectively. |
| 597 | const int num_inputs = inputs.size(); |
| 598 | OperandLabels free_labels(num_inputs); |
| 599 | gtl::InlinedVector<Tensor, 2> inputs_reduced(num_inputs); |
| 600 | gtl::InlinedVector<bool, 2> swap_free_and_contract(num_inputs); |
| 601 | for (int i = 0; i < num_inputs; ++i) { |
| 602 | OP_REQUIRES_OK(ctx, |
| 603 | ReduceOperand<Device, T>( |
| 604 | ctx, inputs[i], label_types, input_label_counts[i], |
| 605 | &input_labels[i], &free_labels[i], |
| 606 | &swap_free_and_contract[i], &inputs_reduced[i])); |
| 607 | } |
| 608 | |
| 609 | // After reduction, the inputs should be reshaped to Tensors suitable for |
| 610 | // contraction. If num_inputs is 1, the reduced input is simply forwarded to |
| 611 | // the output. |
| 612 | Tensor contraction_output_reshaped; |
| 613 | OP_REQUIRES_OK(ctx, ContractOperands<Device, T>( |
| 614 | ctx, inputs_reduced, swap_free_and_contract, |
| 615 | &contraction_output_reshaped)); |
| 616 | |
| 617 | // Copy the batch labels from the contraction output. Recover the batch |
| 618 | // shape, which may have been broadcasted. |
| 619 | TensorShape result_shape = contraction_output_reshaped.shape(); |
| 620 | result_shape.RemoveLastDims(2); |
| 621 | |
| 622 | int num_labels = label_types.size(); |
| 623 | Labels result_labels; |
| 624 | // All batch dimensions should be present in the contracted result. First |
| 625 | // the broadcasting dimensions, then the named batch dimensions. |
| 626 | for (int label = 0; label < num_labels; ++label) { |
| 627 | if (label_types[label] == kBroadcasting) result_labels.push_back(label); |
| 628 | } |
| 629 | for (int label = 0; label < num_labels; ++label) { |
| 630 | if (label_types[label] == kBatch) result_labels.push_back(label); |
| 631 | } |
| 632 | for (int i = 0; i < num_inputs; ++i) { |
nothing calls this directly
no test coverage detected