| 700 | ~BaseBatchMatMulOp() override {} |
| 701 | |
| 702 | void Compute(OpKernelContext* ctx) override { |
| 703 | const Tensor& in0 = ctx->input(0); |
| 704 | const Tensor& in1 = ctx->input(1); |
| 705 | |
| 706 | ValidateInputTensors(ctx, in0, in1); |
| 707 | |
| 708 | MatMulBCast bcast(in0.shape().dim_sizes(), in1.shape().dim_sizes()); |
| 709 | OP_REQUIRES( |
| 710 | ctx, bcast.IsValid(), |
| 711 | errors::InvalidArgument( |
| 712 | "In[0] and In[1] must have compatible batch dimensions: ", |
| 713 | in0.shape().DebugString(), " vs. ", in1.shape().DebugString())); |
| 714 | |
| 715 | TensorShape out_shape = bcast.output_batch_shape(); |
| 716 | auto batch_size = bcast.output_batch_size(); |
| 717 | auto d0 = in0.dim_size(in0.dims() - 2); |
| 718 | auto d1 = in0.dim_size(in0.dims() - 1); |
| 719 | Tensor in0_reshaped; |
| 720 | OP_REQUIRES( |
| 721 | ctx, |
| 722 | in0_reshaped.CopyFrom(in0, TensorShape({bcast.x_batch_size(), d0, d1})), |
| 723 | errors::Internal("Failed to reshape In[0] from ", |
| 724 | in0.shape().DebugString())); |
| 725 | auto d2 = in1.dim_size(in1.dims() - 2); |
| 726 | auto d3 = in1.dim_size(in1.dims() - 1); |
| 727 | Tensor in1_reshaped; |
| 728 | OP_REQUIRES( |
| 729 | ctx, |
| 730 | in1_reshaped.CopyFrom(in1, TensorShape({bcast.y_batch_size(), d2, d3})), |
| 731 | errors::Internal("Failed to reshape In[1] from ", |
| 732 | in1.shape().DebugString())); |
| 733 | if (adj_x_ || trans_x_) std::swap(d0, d1); |
| 734 | if (adj_y_ || trans_y_) std::swap(d2, d3); |
| 735 | OP_REQUIRES(ctx, d1 == d2, |
| 736 | errors::InvalidArgument( |
| 737 | "In[0] mismatch In[1] shape: ", d1, " vs. ", d2, ": ", |
| 738 | in0.shape().DebugString(), " ", in1.shape().DebugString(), |
| 739 | " ", adj_x_, " ", adj_y_)); |
| 740 | out_shape.AddDim(d0); |
| 741 | out_shape.AddDim(d3); |
| 742 | Tensor* out = nullptr; |
| 743 | OP_REQUIRES_OK(ctx, ctx->allocate_output(0, out_shape, &out)); |
| 744 | if (out->NumElements() == 0) { |
| 745 | return; |
| 746 | } |
| 747 | if (in0.NumElements() == 0 || in1.NumElements() == 0) { |
| 748 | functor::SetZeroFunctor<Device, Scalar> f; |
| 749 | f(ctx->eigen_device<Device>(), out->flat<Scalar>()); |
| 750 | return; |
| 751 | } |
| 752 | Tensor out_reshaped; |
| 753 | OP_REQUIRES(ctx, |
| 754 | out_reshaped.CopyFrom(*out, TensorShape({batch_size, d0, d3})), |
| 755 | errors::Internal("Failed to reshape output from ", |
| 756 | out->shape().DebugString())); |
| 757 | |
| 758 | if (std::is_same<Scalar, bfloat16>::value) { |
| 759 | bool is_cpu = std::is_same<Device, CPUDevice>::value; |
nothing calls this directly
no test coverage detected