experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Deal with disabled-jit sooner.

Browse Source 
PiperOrigin-RevId: 360902092
Change-Id: I0eb398b8a8aa10cbf915e6a60aaa93fc34a8c590
This commit is contained in:
Jean-Baptiste Lespiau 2021-03-04 06:30:02 -08:00 committed by TensorFlower Gardener
parent e582ce05be
commit e3c1779bd9
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/compiler/xla/python/jax_jit.cc
View File

@ -1097,6 +1097,9 @@ CacheEntry* CompiledFunction::AddCacheEntry(const py::args& args,
}
py::object CompiledFunction::Call(py::args args, py::kwargs kwargs) {
if (JitIsDisabled()) {
return fun_(*args, **kwargs);
}
if (always_fallback_to_python_) {
return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0];
}
@ -1131,9 +1134,7 @@ py::object CompiledFunction::Call(py::args args, py::kwargs kwargs) {
}
}
CHECK(default_device_);
if (JitIsDisabled()) {
return fun_(*args, **kwargs);
}
ParsedArgumentsAsBuffers arguments;
if (!ParseArguments(args, kwargs, static_argnums_, arguments).ok()) {
return py::cast<py::tuple>(cache_miss_(*args, **kwargs))[0];