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 __anondb58f4d00102() local
38 GetNodeAttr(c->attrs(), "num_computations", &num_computations)); in __anondb58f4d00102()
42 for (int i = 0; i < num_computations; ++i) { in __anondb58f4d00102()
96 int num_computations; in __anondb58f4d00302() local
98 GetNodeAttr(c->attrs(), "num_computations", &num_computations)); in __anondb58f4d00302()
102 for (int i = 0; i < num_computations; ++i) { in __anondb58f4d00302()
106 for (int i = 0; i < num_computations; ++i) { in __anondb58f4d00302()
107 c->set_output(num_computations + i + 1, c->Scalar()); in __anondb58f4d00302()
/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.h70 int num_computations, bool return_hlo_protos, in TpuCompileOpKernelCommon() argument
75 num_computations_(num_computations), in TpuCompileOpKernelCommon()
84 int num_computations, bool return_hlo_protos, in TpuCompileOpKernelCommon() argument
90 num_computations_(num_computations), in TpuCompileOpKernelCommon()
Dtpu_compile_op_support.cc393 int num_computations; in CompileOpMetadataFromContext() local
394 TF_RETURN_IF_ERROR(ctx->GetAttr("num_computations", &num_computations)); in CompileOpMetadataFromContext()
410 if (num_computations != metadata->num_cores_per_replica()) { in CompileOpMetadataFromContext()
415 num_computations, " vs ", metadata->num_cores_per_replica(), ")"); in CompileOpMetadataFromContext()
/external/tensorflow/tensorflow/compiler/xla/python/
Dpy_executable.cc143 int num_computations = executable_->addressable_devices().size(); in ExecuteShardedOnLocalDevices() local
147 if (arg.size() != num_computations) { in ExecuteShardedOnLocalDevices()
151 num_computations, in ExecuteShardedOnLocalDevices()
159 std::vector<std::vector<PjRtBuffer*>> arg_buffers(num_computations); in ExecuteShardedOnLocalDevices()
160 for (int computation = 0; computation < num_computations; ++computation) { in ExecuteShardedOnLocalDevices()
175 outputs[buffer_id].reserve(num_computations); in ExecuteShardedOnLocalDevices()
176 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.pbtxt53735 number_attr: "num_computations"
53740 number_attr: "num_computations"
53743 name: "num_computations"
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/
Dtf_generated_ops.td18543 TF_DerivedResultSizeAttr num_computations = TF_DerivedResultSizeAttr<1>;