Searched refs:jitlib (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | jax_jit.cc | 1210 py::module jitlib = m.def_submodule("jax_jit", "Jax C++ jit library"); in BuildJaxjitSubmodule() local
1213 jitlib, "CompiledFunction"); in BuildJaxjitSubmodule()
1218 jitlib.def("set_disable_jit", &SetDisableJit); in BuildJaxjitSubmodule()
1219 jitlib.def("get_disable_jit", &GetDisableJit); in BuildJaxjitSubmodule()
1220 jitlib.def( in BuildJaxjitSubmodule()
1235 jitlib.def("device_put", [](py::handle obj, bool jax_enable_x64, in BuildJaxjitSubmodule()
1261 py::class_<ArgSignature> arg_signature(jitlib, "ArgSignature"); in BuildJaxjitSubmodule()
1272 jitlib.def("_ArgSignatureOfValue", &ArgSignatureOfValue); in BuildJaxjitSubmodule()
1276 jitlib.def("_DtypeTo32BitDtype", [](const py::object obj) -> py::object { in BuildJaxjitSubmodule()
1285 jitlib.def("_is_float0", &IsFloat0); in BuildJaxjitSubmodule()
[all …]
|