Searched refs:get_jax_enable_x64 (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | pmap_lib.cc | 135 py::function get_jax_enable_x64, std::vector<int> static_argnums) in PmapFunction() argument
139 get_jax_enable_x64_(get_jax_enable_x64) { in PmapFunction()
425 py::function get_jax_enable_x64, in BuildPmapSubmodule()
429 std::move(get_jax_enable_x64), std::move(static_argnums)); in BuildPmapSubmodule()
|
| D | jax_jit.cc | 844 py::function get_device, py::function get_jax_enable_x64,
918 py::function get_jax_enable_x64, in CompiledFunction() argument
924 get_jax_enable_x64_(get_jax_enable_x64), in CompiledFunction()
1223 py::function get_jax_enable_x64, py::function get_jax_disable_jit, in BuildJaxjitSubmodule()
1227 std::move(get_jax_enable_x64), std::move(get_jax_disable_jit), in BuildJaxjitSubmodule()
|