Home
last modified time | relevance | path

Searched refs:execution_input (Results 1 – 7 of 7) sorted by relevance

/external/tensorflow/tensorflow/compiler/xrt/kernels/
Dxrt_execute_op.cc214 xla::ExecutionInput* execution_input = &(*execution_inputs)[i]; in UpdateDynamicInputs() local
225 xla::ShapeUtil::TryGetSubshape(execution_input->shape(), index)); in UpdateDynamicInputs()
235 execution_input->Buffer(index).AsDeviceMemoryBase(); in UpdateDynamicInputs()
244 execution_input->SetBuffer( in UpdateDynamicInputs()
246 execution_input->ClearUnownedIndex(index); in UpdateDynamicInputs()
251 TF_RETURN_IF_ERROR(execution_input->SetDynamicShape(compile_time_shape)); in UpdateDynamicInputs()
253 execution_input->ToShapedBuffer( in UpdateDynamicInputs()
446 const Tensor& execution_input = context->input(0); in DoWork() local
447 TF_RET_CHECK(TensorShapeUtils::IsScalar(execution_input.shape())); in DoWork()
448 int64_t compilation_handle = execution_input.scalar<int64_t>()(); in DoWork()
Dtpu_execute_op.cc295 const Tensor& execution_input = context->input(0); in DoWork() local
296 TF_RET_CHECK(TensorShapeUtils::IsScalar(execution_input.shape())); in DoWork()
297 int64_t compilation_handle = execution_input.scalar<int64_t>()(); in DoWork()
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dtracked_device_buffer.cc153 ExecutionInput* execution_input, in AddToInputAsDonated() argument
160 execution_input->SetUnownedIndex((*iterator)->first); in AddToInputAsDonated()
Dpjrt_stream_executor_client.cc576 ExecutionInput* execution_input, in AddToInput() argument
580 buffer()->AddToInputAsDonated(iterator, end, execution_input, allocator); in AddToInput()
1638 ExecutionInput execution_input; member
1710 ExecutionInput execution_input(tupled_parameter_shape); in MakeTupleHelper() local
1712 execution_input.MutableBuffers()->begin(); in MakeTupleHelper()
1714 execution_input.MutableBuffers()->end(); in MakeTupleHelper()
1716 execution_input.SetBuffer( in MakeTupleHelper()
1723 device_buffer.AddToInput(&input_iterator, iterator_end, &execution_input, in MakeTupleHelper()
1729 stream, execution_input.Buffers())); in MakeTupleHelper()
1739 return TupleHandle({std::move(execution_input), std::move(transfer_event)}); in MakeTupleHelper()
[all …]
Dtracked_device_buffer.h168 ExecutionInput* execution_input,
Dpjrt_stream_executor_client.h511 ExecutionInput* execution_input,
/external/tensorflow/tensorflow/compiler/jit/
Dxla_launch_util.cc227 static void PopulateExecutionInputBuffer(xla::ExecutionInput& execution_input, in PopulateExecutionInputBuffer() argument
233 execution_input.MutableBuffer(index); in PopulateExecutionInputBuffer()
297 xla::ExecutionInput& execution_input = arguments.back(); in PopulateInputs() local
299 PopulateExecutionInputBuffer(execution_input, xla::ShapeIndex{}, dmem, in PopulateInputs()