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

diff --git a/tensorflow/python/pywrap_tfe.i b/tensorflow/python/pywrap_tfe.i
index 7acb8ee..5ee5530 100644 (file)
--- 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);
   }
 }