blob: b3f5703ccf452a1a933ca8904f2be669428379f4 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
|
#ifndef TENSORFLOW_KERNELS_SENDRECV_OPS_H_
#define TENSORFLOW_KERNELS_SENDRECV_OPS_H_
#include "tensorflow/core/framework/op_kernel.h"
namespace tensorflow {
class SendOp : public OpKernel {
public:
explicit SendOp(OpKernelConstruction* ctx);
void Compute(OpKernelContext* ctx) override;
private:
string key_prefix_;
TF_DISALLOW_COPY_AND_ASSIGN(SendOp);
};
class RecvOp : public AsyncOpKernel {
public:
explicit RecvOp(OpKernelConstruction* ctx);
void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override;
private:
string key_prefix_;
TF_DISALLOW_COPY_AND_ASSIGN(RecvOp);
};
} // end namespace tensorflow
#endif // TENSORFLOW_KERNELS_SENDRECV_OPS_H_
|