Searched refs:cache_miss_ (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | pmap_lib.cc | 137 cache_miss_(std::move(cache_miss)), in PmapFunction()
176 const py::function cache_miss_; member in jax::PmapFunction
257 return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0]; in Call()
266 return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0]; in Call()
273 return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0]; in Call()
282 py::object out_and_fastpath_data = cache_miss_(*args, **kwargs); in Call()
300 return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0]; in Call()
|
| D | jax_jit.cc | 878 const py::function cache_miss_; member in jax::__anone04173081411::CompiledFunction
922 cache_miss_(std::move(cache_miss)), in CompiledFunction()
1107 return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0]; in Call()
1128 return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0]; in Call()
1134 return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0]; in Call()
1147 return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0]; in Call()
1155 return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0]; in Call()
1161 py::object out_and_fastpath_data = cache_miss_(*args, **kwargs); in Call()
1178 return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0]; in Call()
|