Searched refs:arg_buffers (Results 1 – 7 of 7) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | py_executable.cc | 95 std::vector<PjRtBuffer*> arg_buffers(args.size()); in Execute() local
96 absl::c_transform(args, arg_buffers.begin(), in Execute()
99 executable_->Execute({arg_buffers}, options_)); in Execute()
117 std::vector<std::vector<PjRtBuffer*>> arg_buffers(args.size()); in ExecuteOnLocalDevices() local
119 arg_buffers[computation].resize(args[computation].size()); in ExecuteOnLocalDevices()
120 absl::c_transform(args[computation], arg_buffers[computation].begin(), in ExecuteOnLocalDevices()
124 executable_->Execute(arg_buffers, options_)); in ExecuteOnLocalDevices()
159 std::vector<std::vector<PjRtBuffer*>> arg_buffers(num_computations); in ExecuteShardedOnLocalDevices() local
161 arg_buffers[computation].resize(args.size()); in ExecuteShardedOnLocalDevices()
162 absl::c_transform(args, arg_buffers[computation].begin(), in ExecuteShardedOnLocalDevices()
[all …]
|
| D | pmap_lib.cc | 315 std::vector<std::vector<xla::PyBuffer*>> arg_buffers; in Call() local
316 arg_buffers.reserve(list_of_list_of_buffers.size()); in Call()
323 arg_buffers.push_back(std::move(buffers)); in Call()
327 ValueOrThrow(cache_entry->executable->ExecuteOnLocalDevices(arg_buffers)); in Call()
|
| D | xla_client.py | 375 arg_buffers = []
377 arg_buffers.append(flat_arg_buffers[:len(replica_args)])
381 for xs in executable.execute_on_local_devices(arg_buffers)]
|
| D | jax_jit.h | 126 std::vector<xla::PjRtBuffer*> arg_buffers; member
|
| D | jax_jit.cc | 945 std::vector<xla::PjRtBuffer*>& arg_buffers = arguments.arg_buffers; in ConvertArgsToBuffers() local
949 arg_buffers.reserve(num_flat_dynamic_args); in ConvertArgsToBuffers()
1013 arg_buffers.push_back(buffer); in ConvertArgsToBuffers()
1181 ValueOrThrow(cache_entry->executable->PjRtExecute(arguments.arg_buffers)); in Call()
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_on_demand_compiler.cc | 110 auto* arg_buffers = arg.MutableBuffers(); in ExecuteAsyncOnStream() local
112 for (auto& pair : *arg_buffers) { in ExecuteAsyncOnStream()
|
| /external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/ |
| D | tpu_client.cc | 773 std::vector<std::vector<PyTpuBuffer*>> arg_buffers(num_computations); in ExecuteShardedOnLocalDevices() local
775 arg_buffers[computation].resize(args.size()); in ExecuteShardedOnLocalDevices()
777 args, arg_buffers[computation].begin(), in ExecuteShardedOnLocalDevices()
780 TF_ASSIGN_OR_RETURN(output_buffers, ExecuteOnLocalDevices(arg_buffers)); in ExecuteShardedOnLocalDevices()
|