Home
last modified time | relevance | path

Searched refs:TpuComputationPlacer (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dtpu_computation_placer.cc27 using StatusOr = TpuComputationPlacer::StatusOr<T>;
29 TpuComputationPlacer::TpuComputationPlacer() { in TpuComputationPlacer() function in tensorflow::tpu::TpuComputationPlacer
33 TpuComputationPlacer::~TpuComputationPlacer() { in ~TpuComputationPlacer()
37 StatusOr<int> TpuComputationPlacer::DeviceId(int replica, int computation, in DeviceId()
43 StatusOr<xla::DeviceAssignment> TpuComputationPlacer::AssignDevices( in AssignDevices()
57 TpuComputationPlacer::AssignLocalDevices(TpuHostLocationExternal host_location, in AssignLocalDevices()
72 return std::make_unique<TpuComputationPlacer>(); in CreateTpuComputationPlacer()
Dtpu_computation_placer.h27 class TpuComputationPlacer : public xla::ComputationPlacer {
32 TpuComputationPlacer();
33 ~TpuComputationPlacer() override;
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dtpu_client.cc107 return tf_tpu::TpuComputationPlacer::AssignLocalDevices(host, num_replicas, in GetDefaultDeviceAssignment()