| /external/tensorflow/tensorflow/core/common_runtime/ |
| D | shape_refiner.cc | 71 InferenceContext* node_context = CHECK_NOTNULL(GetContext(node)); in InferShapesForFunctionSubNode() local
92 node_context->set_output(0, node_context->UnknownShape()); in InferShapesForFunctionSubNode()
94 node_context->set_output(0, outer_context->input(index)); in InferShapesForFunctionSubNode()
99 node_context->set_output_handle_shapes_and_types(0, *resource); in InferShapesForFunctionSubNode()
119 node_context->ShapeHandleToProto(node_context->input(0), &proto); in InferShapesForFunctionSubNode()
123 auto* resource = node_context->input_handle_shapes_and_types(0); in InferShapesForFunctionSubNode()
297 InferenceContext* node_context = node_ext_context->get_context(); in UpdateNode() local
300 TF_RETURN_IF_ERROR(node_context->construction_status()); in UpdateNode()
320 ShapeHandle existing_input = node_context->input(dst_input); in UpdateNode()
322 if (node_context->MergeInput(dst_input, c->output(src_output))) { in UpdateNode()
[all …]
|
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | tpu_execute_op.cc | 138 tpu::TpuNodeContext* node_context, se::Stream* stream, in AllocateOutputTuple() argument
147 node_context->backend()->transfer_manager()->ReadDynamicShapes( in AllocateOutputTuple()
157 node_context->backend(), device_ordinal, &output_tuple)); in AllocateOutputTuple()
172 tpu::TpuNodeContext* node_context, se::Stream* stream, in AllocateOutputTensors() argument
180 AllocateOutputTuple(node_context, stream, input_tuples, in AllocateOutputTensors()
188 OpKernelContext* context, tpu::TpuNodeContext* node_context, in RunExecutable() argument
219 node_context, device_assignment.get(), in RunExecutable()
230 OpKernelContext* context, tpu::TpuNodeContext* node_context, in ExecuteTPUProgram() argument
238 return RunExecutable(context, node_context, executable, in ExecuteTPUProgram()
276 TF_ASSIGN_OR_RETURN(std::unique_ptr<tpu::TpuNodeContext> node_context, in DoWork()
[all …]
|
| /external/tensorflow/tensorflow/stream_executor/tpu/ |
| D | tpu_node_context.h | 50 explicit TpuNodeContext(int device_ordinal, XLA_TpuNodeContext* node_context) in TpuNodeContext() argument
51 : device_ordinal_(device_ordinal), node_context_(node_context) { in TpuNodeContext()
52 CHECK_NE(node_context, nullptr); in TpuNodeContext()
|
| D | tpu_node_context.cc | 32 XLA_TpuNodeContext* node_context = in Create() local
37 tpu::OpsApiFn()->TpuNodeContext_FreeFn(node_context); in Create()
40 return std::make_unique<TpuNodeContext>(device_ordinal, node_context); in Create()
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_execute.cc | 418 TpuNodeContext* node_context, xla::DeviceAssignment* device_assignment, in TPUExecute() argument
426 VLOG(1) << "TPUExecute on device " << node_context->device_ordinal(); in TPUExecute()
428 xla::Backend* backend = node_context->backend(); in TPUExecute()
432 std::make_shared<HostTransferManager>(node_context, backend); in TPUExecute()
438 << node_context->device_ordinal(); in TPUExecute()
484 const int32 device_ordinal = node_context->device_ordinal(); in TPUExecute()
|
| D | tpu_execute.h | 46 tpu::TpuNodeContext* node_context, xla::DeviceAssignment* device_assignment,
|
| D | tpu_ops_c_api.h | 421 void TpuNodeContext_Free(XLA_TpuNodeContext* node_context);
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_execute_op.cc | 397 const VariableUpdateMap& variable_updates, TpuNodeContext* node_context, in AllocateOutputTensors() argument
413 node_context->backend()->transfer_manager(); in AllocateOutputTensors()
437 node_context->backend()->memory_allocator(); in AllocateOutputTensors()
628 TF_ASSIGN_OR_RETURN(std::unique_ptr<TpuNodeContext> node_context, in DoWork()
656 xla::Backend* const backend = node_context->backend(); in DoWork()
758 rendezvous_key_base, rng_seed, node_context.get(), in DoWork()
768 variable_update_map, node_context.get(), stream, device_ordinal, in DoWork()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/translate/ |
| D | import_model.cc | 792 auto* node_context = shape_refiner_->GetContext(node); in AddNodesToShapeRefiner() local
795 node_context->MakeShapeFromShapeProto(shape.value(), &handle); in AddNodesToShapeRefiner()
799 node_context->set_output(shape.index(), handle); in AddNodesToShapeRefiner()
827 auto* node_context = shape_refiner_->GetContext(node); in AddNodesToShapeRefiner() local
828 DCHECK(node_context != nullptr); in AddNodesToShapeRefiner()
832 node_context->MakeShapeFromShapeProto(attr->shape(), &handle); in AddNodesToShapeRefiner()
836 node_context->set_output(0, handle); in AddNodesToShapeRefiner()
840 node_context->set_output(0, node_context->UnknownShape()); in AddNodesToShapeRefiner()
|