Searched refs:PmapFunction (Results 1 – 1 of 1) sorted by relevance
132 class PmapFunction { class134 PmapFunction(py::function fun, py::function cache_miss, in PmapFunction() function in jax::PmapFunction143 ~PmapFunction() { in ~PmapFunction()192 PmapCacheEntry* PmapFunction::GetCacheEntryIfPresent( in GetCacheEntryIfPresent()209 PmapCacheEntry* PmapFunction::AddCacheEntry(const py::args& args, in AddCacheEntry()255 py::object PmapFunction::Call(py::args args, py::kwargs kwargs) { in Call()417 py::class_<PmapFunction, std::unique_ptr<PmapFunction>> cfun(pmap_lib, in BuildPmapSubmodule()419 cfun.def("__call__", &PmapFunction::Call); in BuildPmapSubmodule()420 cfun.def_property_readonly("__signature__", &PmapFunction::PythonSignature); in BuildPmapSubmodule()426 std::vector<int> static_argnums) -> std::unique_ptr<PmapFunction> { in BuildPmapSubmodule()[all …]