Home
last modified time | relevance | path

Searched refs:num_computations (Results 1 – 10 of 10) sorted by relevance

/external/tensorflow/tensorflow/core/tpu/ops/
Dtpu_compile_op.cc36 int num_computations; in __anon78c361350102() local
37 TF_RETURN_IF_ERROR(c->GetAttr("num_computations", &num_computations)); in __anon78c361350102()
41 for (int i = 0; i < num_computations; ++i) { in __anon78c361350102()
95 int num_computations; in __anon78c361350302() local
96 TF_RETURN_IF_ERROR(c->GetAttr("num_computations", &num_computations)); in __anon78c361350302()
100 for (int i = 0; i < num_computations; ++i) { in __anon78c361350302()
104 for (int i = 0; i < num_computations; ++i) { in __anon78c361350302()
105 c->set_output(num_computations + i + 1, c->Scalar()); in __anon78c361350302()
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_compile_op_impl.h40 int num_computations, bool return_hlo_protos, in TpuCompileOpKernelImpl() argument
42 : TpuCompileOpKernelCommon(mlir_module, metadata, num_computations, in TpuCompileOpKernelImpl()
48 int num_computations, bool return_hlo_protos, in TpuCompileOpKernelImpl() argument
51 function, metadata, num_computations, return_hlo_protos, in TpuCompileOpKernelImpl()
Dtpu_compile_op_common.h68 int num_computations, bool return_hlo_protos, in TpuCompileOpKernelCommon() argument
73 num_computations_(num_computations), in TpuCompileOpKernelCommon()
82 int num_computations, bool return_hlo_protos, in TpuCompileOpKernelCommon() argument
88 num_computations_(num_computations), in TpuCompileOpKernelCommon()
Dtpu_compile_op_support.cc399 int num_computations; in CompileOpMetadataFromContext() local
400 TF_RETURN_IF_ERROR(ctx->GetAttr("num_computations", &num_computations)); in CompileOpMetadataFromContext()
416 if (num_computations != metadata->num_cores_per_replica()) { in CompileOpMetadataFromContext()
421 num_computations, " vs ", metadata->num_cores_per_replica(), ")"); in CompileOpMetadataFromContext()
/external/tensorflow/tensorflow/compiler/xla/python/
Dpy_executable.cc95 int num_computations = executable_->addressable_devices().size(); in ExecuteShardedOnLocalDevices() local
99 if (arg.size() != num_computations) { in ExecuteShardedOnLocalDevices()
103 num_computations, in ExecuteShardedOnLocalDevices()
111 std::vector<std::vector<PjRtBuffer*>> arg_buffers(num_computations); in ExecuteShardedOnLocalDevices()
113 for (int computation = 0; computation < num_computations; ++computation) { in ExecuteShardedOnLocalDevices()
128 outputs[buffer_id].reserve(num_computations); in ExecuteShardedOnLocalDevices()
129 for (int computation = 0; computation < num_computations; ++computation) { in ExecuteShardedOnLocalDevices()
/external/tensorflow/tensorflow/core/ops/compat/ops_history_v2/
DTPUCompile.pbtxt19 number_attr: "num_computations"
24 number_attr: "num_computations"
27 name: "num_computations"
/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
Dtpu_client.cc759 int num_computations = args.front().size(); in ExecuteShardedOnLocalDevices() local
761 if (arg.size() != num_computations) { in ExecuteShardedOnLocalDevices()
765 num_computations, in ExecuteShardedOnLocalDevices()
773 std::vector<std::vector<PyTpuBuffer*>> arg_buffers(num_computations); in ExecuteShardedOnLocalDevices()
774 for (int computation = 0; computation < num_computations; ++computation) { in ExecuteShardedOnLocalDevices()
785 outputs[buffer_id].reserve(num_computations); in ExecuteShardedOnLocalDevices()
786 for (int computation = 0; computation < num_computations; ++computation) { in ExecuteShardedOnLocalDevices()
/external/tensorflow/tensorflow/core/api_def/base_api/
Dapi_def_TPUCompile.pbtxt7 'num_computations' is the number of computations to be compiled.
/external/tensorflow/tensorflow/core/ops/
Dops.pbtxt53341 number_attr: "num_computations"
53346 number_attr: "num_computations"
53349 name: "num_computations"
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/
Dtf_generated_ops.td20073 TF_DerivedResultSizeAttr num_computations = TF_DerivedResultSizeAttr<1>;