Searched refs:pmap_data (Results 1 – 1 of 1) sorted by relevance
229 py::dict pmap_data = py::cast<py::dict>(tuple[1]); in AddCacheEntry() local230 if (py::cast<int>(pmap_data["version"]) != 1) { in AddCacheEntry()234 py::cast<int>(pmap_data["version"]), in AddCacheEntry()236 py::cast<std::string>(py::str(py::repr(pmap_data))))); in AddCacheEntry()244 py::cast<std::shared_ptr<xla::PyExecutable>>(pmap_data["xla_executable"]); in AddCacheEntry()246 cache_entry->handle_args = py::cast<py::function>(pmap_data["in_handler"]); in AddCacheEntry()247 cache_entry->out_handler = py::cast<py::function>(pmap_data["out_handler"]); in AddCacheEntry()248 auto out_tree = py::cast<xla::PyTreeDef>(pmap_data["out_pytree_def"]); in AddCacheEntry()