Allow TFE_NewContext to fail more reasonably when SWIG is checking status.

Before:
TFE_Context would check nullptr, and the function would fail straight away.

Now:
TFE_Context is nullptr, so it skips down to checking the status, and an error
is raised.

I'm not able to find in SWIG documentation how to order typemaps in the
generated code - ideally, I'd order it to check the status typemap first. This
code makes it not dependent on this ordering either way.

PiperOrigin-RevId: 191905893

1 files changed, 3 insertions, 3 deletions

diff --git a/tensorflow/python/pywrap_tfe.i b/tensorflow/python/pywrap_tfe.i
index 7acb8eeb1a..5ee55301df 100644
--- a/tensorflow/python/pywrap_tfe.i
+++ b/tensorflow/python/pywrap_tfe.i
@@ -120,9 +120,9 @@ limitations under the License.
 
 }
 %typemap(out) (TFE_Context*) {
-  if ($1 == nullptr) {
-    SWIG_fail;
-  } else {
+  // When the TFE_Context* returned is a nullptr, we expect the status is not
+  // OK. This will raise an error (happens in another typemap).
+  if ($1 != nullptr) {
     $result = PyCapsule_New($1, nullptr, TFE_DeleteContextCapsule);
   }
 }