| 84 | } |
| 85 | |
| 86 | void Compile(XlaOpKernelContext* ctx) override { |
| 87 | OP_REQUIRES( |
| 88 | ctx, ctx->num_inputs() == 2, |
| 89 | errors::Unimplemented("Broadcast for n-ary operations (n > 2)")); |
| 90 | |
| 91 | absl::InlinedVector<BCast::Vec, 4> shapes; |
| 92 | for (int i = 0; i < ctx->num_inputs(); ++i) { |
| 93 | const TensorShape in_shape = ctx->InputShape(i); |
| 94 | OP_REQUIRES(ctx, TensorShapeUtils::IsVector(in_shape), |
| 95 | errors::InvalidArgument("In[", i, "] must be a vector.", |
| 96 | in_shape.DebugString())); |
| 97 | std::vector<int64> vec; |
| 98 | OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(i, &vec)); |
| 99 | |
| 100 | shapes.push_back(BCast::Vec(vec.begin(), vec.end())); |
| 101 | } |
| 102 | BCast bcast(shapes[0], shapes[1]); |
| 103 | OP_REQUIRES(ctx, bcast.IsValid(), |
| 104 | errors::InvalidArgument( |
| 105 | "Incompatible shapes: [", absl::StrJoin(shapes[0], ","), |
| 106 | "] vs. [", absl::StrJoin(shapes[1], ","), "]")); |
| 107 | Output(ctx, 0, bcast.grad_x_reduce_idx()); |
| 108 | Output(ctx, 1, bcast.grad_y_reduce_idx()); |
| 109 | } |
| 110 | |
| 111 | private: |
| 112 | void Output(XlaOpKernelContext* ctx, int idx, const BCast::Vec& v) { |
nothing calls this directly
no test coverage detected