| 95 | } |
| 96 | |
| 97 | Status ComputeBinOp(OpKernelContext* op_ctx, OpKernelContext::Params* params, |
| 98 | Device* device, OpKernel* op, Tensor* output, |
| 99 | Tensor* input) { |
| 100 | // Prepare an OpKernelContext that is identical to that of the original Op |
| 101 | // (i.e. the collective), except for the input output sizes and identities and |
| 102 | // the Op itself. |
| 103 | // TODO(ayushd, tucker): Is it possible to cache and reuse these objects? |
| 104 | // They're mostly identical inside one device execution. |
| 105 | std::unique_ptr<SubContext> sub_ctx( |
| 106 | new SubContext(op_ctx, params, op, output, input)); |
| 107 | device->Compute(op, sub_ctx->sub_ctx_.get()); |
| 108 | return sub_ctx->sub_ctx_->status(); |
| 109 | } |
| 110 | |
| 111 | } // namespace collective_util |
| 112 | } // namespace tensorflow |
no test coverage detected