Searched defs:tpu_core_ids (Results 1 – 5 of 5) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/python/ |
| H A D | dtensor_device.py | 328 def set_tpu_core_ids(self, mesh_name, tpu_core_ids): argument
341 def tpu_core_ids_to_locations(self, tpu_core_ids): argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/ |
| H A D | pywrap_dtensor_device.cc | 175 const std::vector<int>& tpu_core_ids) { in PYBIND11_MODULE()
193 const std::vector<int>& tpu_core_ids) { in PYBIND11_MODULE()
|
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/cc/ |
| H A D | dtensor_device.cc | 162 const std::vector<int>& tpu_core_ids) { in SetTPUCoreIDs()
184 TFE_Context* context, const std::vector<int>& tpu_core_ids) { in TPUCoreIDsToLocations()
222 std::vector<int> tpu_core_ids; in TPUCoreLocationsToIDs() local
2091 const std::vector<int>& tpu_core_ids, void* device_info, in SetTPUCoreIDs()
2104 TFE_Context* context, const std::vector<int>& tpu_core_ids, in TPUCoreIDsToLocations()
|
| H A D | tensor_layout.cc | 175 std::map<std::string, std::vector<int>>& Mesh::tpu_core_ids() { in tpu_core_ids() function in tensorflow::dtensor::Mesh
176 static auto tpu_core_ids = new std::map<std::string, std::vector<int>>(); in tpu_core_ids() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/utils/ |
| H A D | update_tpu_metadata.cc | 121 const std::vector<int>& tpu_core_ids = Mesh::tpu_core_ids()[mesh_name]; in UpdateTPUCompileMetadata() local
|