Home
last modified time | relevance | path

Searched refs:computation_count (Results 1 – 25 of 32) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dcomputation_placer.cc49 for (int c = 0; c < computation_count(); ++c) { in LogicalIdForDevice()
77 proto->set_computation_count(computation_count()); in Serialize()
78 for (int computation = 0; computation < computation_count(); ++computation) { in Serialize()
90 TF_RET_CHECK(proto.computation_devices_size() == proto.computation_count()); in Deserialize()
91 if (proto.replica_count() <= 0 || proto.computation_count() <= 0) { in Deserialize()
95 proto.replica_count(), proto.computation_count()); in Deserialize()
98 proto.replica_count(), proto.computation_count()); in Deserialize()
99 for (int computation = 0; computation < proto.computation_count(); in Deserialize()
113 string output = StrCat("Computations: ", computation_count(), in ToString()
115 for (int computation = 0; computation < computation_count(); ++computation) { in ToString()
[all …]
Dtpu_computation_placer.cc39 int computation_count) { in DeviceId() argument
44 int replica_count, int computation_count) { in AssignDevices() argument
46 xla::DeviceAssignment result(replica_count, computation_count); in AssignDevices()
48 placer_, replica_count, computation_count, result.data(), in AssignDevices()
59 int computation_count) { in AssignLocalDevices() argument
61 xla::DeviceAssignment result(replica_count, computation_count); in AssignLocalDevices()
63 host_location.impl(), replica_count, computation_count, result.data(), in AssignLocalDevices()
Dcomputation_placer.h44 DeviceAssignment(int replica_count, int computation_count) in DeviceAssignment() argument
45 : Array2D<int>(replica_count, computation_count, -1) { in DeviceAssignment()
47 CHECK_GT(computation_count, 0); in DeviceAssignment()
51 int computation_count() const { return width(); } in computation_count() function
87 int replica_count, int computation_count);
92 int computation_count);
Dhlo_subcomputation_unification_test.cc87 EXPECT_EQ(3, module->computation_count()); in TEST_F()
90 EXPECT_EQ(2, module->computation_count()); in TEST_F()
116 EXPECT_EQ(3, module->computation_count()); in TEST_F()
119 EXPECT_EQ(2, module->computation_count()); in TEST_F()
146 EXPECT_EQ(3, module->computation_count()); in TEST_F()
149 EXPECT_EQ(3, module->computation_count()); in TEST_F()
173 EXPECT_EQ(1, module->computation_count()); in TEST_F()
Dtpu_computation_placer.h36 int computation_count) override;
39 int computation_count) override;
43 int computation_count);
Dhlo_execution_profile.h65 size_t computation_count() const { in computation_count() function
74 return instruction_count() + computation_count() + extra_metrics_count(); in total_count()
Dflatten_call_graph_test.cc218 EXPECT_EQ(7, module->computation_count()); in TEST_F()
245 EXPECT_EQ(2, module->computation_count()); in TEST_F()
251 EXPECT_EQ(3, module->computation_count()); in TEST_F()
Dhlo_schedule_test.cc330 ASSERT_EQ(module->computation_count(), 3); in TEST_F()
332 ASSERT_EQ(module->computation_count(), 1); in TEST_F()
Dcollective_ops_utils.cc133 int partition_count = device_assignment.computation_count(); in GetParticipatingDevicesGroups()
242 int partition_count = device_assignment.computation_count(); in GetParticipatingDevices()
Dcompile_only_service.cc87 options.static_device_assignment().computation_count()); in CompileAheadOfTime()
Dhlo_execution_profile.cc63 hlo_profile_index_map.computation_count()); in CreateHloProfilePrinterData()
Dhlo_module.h184 int64 computation_count() const { return computations_.size(); } in computation_count() function
Dhlo_cost_analysis.cc98 auto computation_count = ShapeUtil::ElementsIn(shape); in HandleElementwiseOp() local
111 current_properties_[kTranscendentalsKey] = computation_count; in HandleElementwiseOp()
115 current_properties_[kFlopsKey] = computation_count; in HandleElementwiseOp()
Dhlo_runner.cc209 int64_t num_partitions = device_assignment->computation_count(); in ExecuteReplicatedImpl()
Dhlo_module.cc469 CHECK_EQ(module_config.static_device_assignment().computation_count(), in CreateModuleConfigFromShape()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dverified_hlo_module.cc30 TF_RET_CHECK(computation_count() == 0); in ParseHloStringAndVerifyModule()
48 if (computation_count() == 0) { in Verify()
/external/tensorflow/tensorflow/compiler/xla/client/
Dlocal_client.cc376 if (options.device_assignment().computation_count() != in Compile()
381 options.device_assignment().computation_count(), in Compile()
/external/tensorflow/tensorflow/stream_executor/tpu/
Dtpu_executor_c_api.h246 int computation_count, int* assignment,
250 int computation_count,
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_compile_op_support.cc436 if (device_assignment.computation_count() != in CompileOpMetadataFromContext()
440 device_assignment.computation_count(), " vs ", in CompileOpMetadataFromContext()
/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
Dtpu_client.cc520 const int num_partitions = device_assignment_.computation_count(); in PyTpuExecutable()
821 } else if (options.device_assignment().computation_count() != 1) { in Compile()
824 options.device_assignment().computation_count()); in Compile()
Dtpu_client.h323 int num_partitions() const { return device_assignment_.computation_count(); } in num_partitions()
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dutils.cc139 *num_partitions = build_options->device_assignment().computation_count(); in ParseDeviceAssignmentCompileOptions()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
Dtpu_rewrite_device_util_test.cc477 EXPECT_EQ(xla_device_assignment->computation_count(), 2); in TEST()
583 EXPECT_EQ(xla_device_assignment->computation_count(), 3); in TEST()
/external/tensorflow/tensorflow/compiler/xla/python/xla_extension/
D__init__.pyi182 def computation_count(self) -> int: ...
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla_compiler.cc501 .def("computation_count", &DeviceAssignment::computation_count) in BuildXlaCompilerSubmodule()

12