Searched defs:python_shard_arg_fallback (Results 1 – 1 of 1) sorted by relevance
197 py::function python_shard_arg_fallback) in PmapFunction()228 const py::function& python_shard_arg_fallback() const { in python_shard_arg_fallback() function in jax::PmapFunction255 py::function fun, cache_miss, python_shard_arg_fallback; in ClearPythonReferences() local711 py::function python_shard_arg_fallback) { in InitializePmapFunction()721 py::function python_shard_arg_fallback) { in MakePmapFunction()903 py::function python_shard_arg_fallback = in BuildPmapSubmodule() local952 py::function python_shard_arg_fallback) -> py::object { in BuildPmapSubmodule()