Searched refs:TpuCoreLocationExternal (Results 1 – 8 of 8) sorted by relevance
23 TpuDimensionsExternal TpuCoreLocationExternal::chip_coordinates() const { in chip_coordinates()30 TpuDimensionsExternal TpuCoreLocationExternal::host_coordinates() const { in host_coordinates()37 int32 TpuCoreLocationExternal::index() const { in index()41 int32 TpuCoreLocationExternal::Id() const { in Id()49 std::vector<TpuCoreLocationExternal> TpuHostLocationExternal::Cores( in Cores()56 std::vector<TpuCoreLocationExternal> result; in Cores()94 TpuCoreLocationExternal TpuTopologyExternal::CoreForId( in CoreForId()96 return TpuCoreLocationExternal( in CoreForId()100 TpuCoreLocationExternal TpuTopologyExternal::Core(TpuCoreTypeEnum core_type, in Core()103 return TpuCoreLocationExternal(tpu::ExecutorApiFn()->TpuTopology_CoreFn( in Core()[all …]
33 class TpuCoreLocationExternal {35 TpuCoreLocationExternal() : core_location_(nullptr) {} in TpuCoreLocationExternal() function36 explicit TpuCoreLocationExternal(SE_TpuTopology_Core* core_location) in TpuCoreLocationExternal() function54 std::vector<TpuCoreLocationExternal> Cores(TpuCoreTypeEnum core_type) const;78 TpuCoreLocationExternal CoreForId(TpuCoreTypeEnum core_type, int id) const;79 TpuCoreLocationExternal Core(TpuCoreTypeEnum core_type, int x, int y, int z,81 std::vector<TpuCoreLocationExternal> cores(TpuCoreTypeEnum core_type) const;
58 virtual TpuCoreLocationExternal GetCoreLocationExternal() const { in GetCoreLocationExternal()
84 tpu::TpuCoreLocationExternal TpuExecutor::GetCoreLocationExternal() const { in GetCoreLocationExternal()85 return tpu::TpuCoreLocationExternal( in GetCoreLocationExternal()
105 tpu::TpuCoreLocationExternal GetCoreLocationExternal() const override;
31 PjRtTpuDevice(const tensorflow::tpu::TpuCoreLocationExternal core, in PjRtTpuDevice()42 const tensorflow::tpu::TpuCoreLocationExternal core() const { return core_; } in core()50 const tensorflow::tpu::TpuCoreLocationExternal core_;
153 for (const tf_tpu::TpuCoreLocationExternal& core : in GetTpuDevices()
1445 xla::Array2D<tpu::TpuCoreLocationExternal>* device_assignment) { in ParseDeviceAssignmentAttr()1463 *device_assignment = xla::Array2D<tpu::TpuCoreLocationExternal>( in ParseDeviceAssignmentAttr()1484 tpu::TpuCoreLocationExternal core_location = in ParseDeviceAssignmentAttr()1526 const xla::Array2D<tpu::TpuCoreLocationExternal>& device_assignment, in BuildGeneralDeviceAssignment()1537 const tpu::TpuCoreLocationExternal& core_location = in BuildGeneralDeviceAssignment()1651 xla::Array2D<tpu::TpuCoreLocationExternal> device_assignment; in BuildDeviceAssignment()