Home
last modified time | relevance | path

Searched refs:tpu_topology (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/stream_executor/tpu/
Dtpu_executor_c_api.h243 int TpuTopology_LogicalDevicesPerHost(SE_TpuTopology* tpu_topology,
245 int TpuTopology_LogicalDevicesPerChip(SE_TpuTopology* tpu_topology,
247 int TpuTopology_HostCount(SE_TpuTopology* tpu_topology);
248 int TpuTopology_ChipsPerHost(SE_TpuTopology* tpu_topology);
250 int TpuTopology_ChipBounds_X(SE_TpuTopology* tpu_topology);
251 int TpuTopology_ChipBounds_Y(SE_TpuTopology* tpu_topology);
252 int TpuTopology_ChipBounds_Z(SE_TpuTopology* tpu_topology);
253 bool TpuTopology_HasChip(SE_TpuTopology* tpu_topology, int x, int y, int z);
254 SE_TpuTopology_Core* TpuTopology_CoreForId(SE_TpuTopology* tpu_topology,
257 SE_TpuTopology_Core* TpuTopology_Core(SE_TpuTopology* tpu_topology,
[all …]
DBUILD365 srcs = ["tpu_topology.cc"],
366 hdrs = ["tpu_topology.h"],
/external/tensorflow/tensorflow/python/tpu/
Dtpu_strategy_util.py136 tpu_topology = topology.Topology(serialized=serialized_topology)
137 _INITIALIZED_TPU_SYSTEMS[tpu_name] = tpu_topology
139 return tpu_topology
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_configuration_ops.cc374 auto tpu_topology = ctx->input(0).scalar<tstring>()(); in Compute() local
377 tpu::OpsApiFn()->SetGlobalTPUArrayOp_DoWorkFn(tpu_topology.size(), in Compute()
378 tpu_topology.data(), status); in Compute()
/external/tensorflow/tensorflow/core/tpu/graph_rewrite/
Ddistributed_tpu_rewrite_pass.cc1385 const tpu::TpuTopologyExternal& tpu_topology, in ParseTopologyAttr() argument
1410 int devices_per_chip = tpu_topology.LogicalDevicesPerChip(kTensorCore); in ParseTopologyAttr()
1412 tpu_topology.chip_bounds().x, tpu_topology.chip_bounds().y, in ParseTopologyAttr()
1413 tpu_topology.chip_bounds().z, devices_per_chip, {-1, -1}); in ParseTopologyAttr()
1422 if (!tpu_topology.HasChip(x, y, z) || core < 0 || in ParseTopologyAttr()
1443 const tpu::TpuTopologyExternal& tpu_topology, int num_replicas, in ParseDeviceAssignmentAttr() argument
1465 int devices_per_chip = tpu_topology.LogicalDevicesPerChip(kTensorCore); in ParseDeviceAssignmentAttr()
1467 tpu_topology.chip_bounds().x, tpu_topology.chip_bounds().y, in ParseDeviceAssignmentAttr()
1468 tpu_topology.chip_bounds().z, devices_per_chip, -1); in ParseDeviceAssignmentAttr()
1478 if (!tpu_topology.HasChip(x, y, z) || core < 0 || in ParseDeviceAssignmentAttr()
[all …]
/external/tensorflow/tensorflow/core/tpu/
Dtpu_ops_c_api.h264 const size_t tpu_topology_size, const char* tpu_topology,