Searched refs:num_devices_per_host (Results 1 – 3 of 3) sorted by relevance
87 std::vector<int32_t> num_devices_per_host; in ConstructDevicesPerHost() local105 num_devices_per_host.push_back(input_tensor.scalar<int32_t>()()); in ConstructDevicesPerHost()107 return num_devices_per_host; in ConstructDevicesPerHost()114 xla::StatusOr<std::vector<int32_t>> num_devices_per_host = in Compute() local116 OP_REQUIRES_OK(ctx, num_devices_per_host.status()); in Compute()127 ctx, ConstructTpuPodState(rmgr, *num_devices_per_host, compilation_cache, in Compute()149 size_t num_devices_per_host = -1; in Compute() local170 if (num_devices_per_host == -1) { in Compute()171 num_devices_per_host = in Compute()175 num_devices_per_host == in Compute()[all …]
146 ResourceMgr* rmgr, const std::vector<int32_t>& num_devices_per_host, in ConstructTpuPodState() argument166 params.num_cores_per_host_size = num_devices_per_host.size(); in ConstructTpuPodState()167 params.num_cores_per_host = num_devices_per_host.data(); in ConstructTpuPodState()
56 ResourceMgr* rmgr, const std::vector<int32_t>& num_devices_per_host,