Searched refs:XlaContext (Results 1 – 14 of 14) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | xla_context.cc | 39 const char XlaContext::kXlaContextResourceName[] = "_xla_context";
43 /* static */ XlaContext& XlaContext::Get(const OpKernelContext* ctx) { in Get()
47 XlaContext* context; in Get()
57 /* static */ XlaContext& XlaContext::Get(const XlaOpKernelContext* ctx) { in Get()
61 void XlaContext::set_args(std::vector<XlaExpression> args) { in set_args()
65 XlaContext::XlaContext( in XlaContext() function in tensorflow::XlaContext
76 string XlaContext::DebugString() { return "TLA JIT context"; } in DebugString()
80 void XlaContext::AddRetval(int retval_index, DataType type, in AddRetval()
90 Status XlaContext::AddConstRetval(int retval_index, DataType dtype, in AddConstRetval()
107 xla::ComputationBuilder* XlaContext::builder() { return builder_; } in builder()
[all …]
|
| D | xla_context.h | 39 class XlaContext : public ResourceBase {
42 static XlaContext& Get(const OpKernelContext* ctx);
43 static XlaContext& Get(const XlaOpKernelContext* ctx);
46 XlaContext(XlaCompiler* compiler, xla::ComputationBuilder* builder,
173 TF_DISALLOW_COPY_AND_ASSIGN(XlaContext);
|
| D | xla_op_kernel.cc | 34 return XlaContext::Get(this).builder(); in builder()
309 XlaContext& xla_context = XlaContext::Get(context_); in ReadVariableInput()
432 XlaContext& xla_context = XlaContext::Get(context_); in AssignVariable()
442 return XlaContext::Get(context_).compiler(); in compiler()
462 return XlaContext::Get(context_).GetOrCreateMax(type); in GetOrCreateMax()
467 return XlaContext::Get(context_).GetOrCreateMin(type); in GetOrCreateMin()
472 return XlaContext::Get(context_).GetOrCreateAdd(type); in GetOrCreateAdd()
477 return XlaContext::Get(context_).GetOrCreateMul(type); in GetOrCreateMul()
|
| D | graph_compiler.h | 58 GraphCompiler(XlaContext* xla_context, XlaCompilationDevice* device, in GraphCompiler()
85 XlaContext* xla_context_;
|
| D | xla_compiler.h | 33 class XlaContext; variable
316 XlaContext* context, std::vector<int>* arg_cores,
|
| D | xla_compiler.cc | 295 Status ExecuteGraph(XlaContext* xla_context, std::unique_ptr<Graph> graph, in ExecuteGraph()
312 step_container->name(), XlaContext::kXlaContextResourceName, in ExecuteGraph()
429 bool use_tuple_arg, xla::ComputationBuilder* builder, XlaContext* context, in BuildArguments()
607 XlaContext* context = in CompileGraph()
608 new XlaContext(this, &builder, options_.allow_cpu_custom_calls, in CompileGraph()
|
| D | graph_compiler.cc | 223 XlaContext& context = XlaContext::Get(op_context); in CompileFunctionalNode()
|
| D | xla_compilation_device.cc | 96 auto* b = XlaContext::Get(context).builder(); in Compute()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | retval_op.cc | 57 XlaContext& tc = XlaContext::Get(ctx); in Compile()
|
| D | sendrecv_ops.cc | 45 XlaCompiler* compiler = XlaContext::Get(ctx).compiler(); in Compile()
76 XlaCompiler* compiler = XlaContext::Get(ctx).compiler(); in Compile()
|
| D | stack_ops.cc | 115 XlaContext& xc = XlaContext::Get(ctx); in Compile()
|
| D | arg_op.cc | 51 const XlaExpression& arg = XlaContext::Get(ctx).args()[index_]; in Compile()
|
| D | index_ops_cpu.cc | 71 OP_REQUIRES(ctx, XlaContext::Get(ctx).allow_cpu_custom_calls(), in Compile()
|
| D | tensor_array_ops.cc | 173 XlaContext& xc = XlaContext::Get(ctx); in Compile()
|