Fix segfault when using unusally shaped tensors.

If a numpy array is created with a shape such that one element is zero and the others sum to a large number, an error will be raised. E.g. the following raises an error:

np.ones((0, 2**31, 2**31))

TF previously would not check the return value of PyArray_SimpleNewFromData, which returns null on such shapes. Now the return value is checked.

PiperOrigin-RevId: 477751929

tensorflow/python/eager/tensor_test.py

@@ -498,6 +498,16 @@ def testNumpyTooManyDimensions(self):
         "can have at most 32 dimensions"):
       t.numpy()
 
+  def testNumpyDimsTooBig(self):
+    # Creating a Numpy array fails in some cases if the product of non-zero
+    # dimensions is very big, even if the shape also has a zero in it.
+    t = array_ops.ones((0, 2**31, 2**31))
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        r"Failed to create numpy array from tensor of shape "
+        r"\[0, 2147483648, 2147483648\]. Numpy error.*array is too big"):
+      t.numpy()
+
 
 class TFETensorUtilTest(test_util.TensorFlowTestCase):
 

tensorflow/python/lib/core/BUILD

@@ -82,6 +82,7 @@ cc_library(
     deps = [
         ":bfloat16_lib",
         ":numpy_lib",
+        ":py_util",
         "//tensorflow/c:c_api_no_xla",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",

tensorflow/python/lib/core/ndarray_tensor_bridge.cc

@@ -13,16 +13,20 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+// clang-format off
 // Must be included first.
 #include "tensorflow/python/lib/core/numpy.h"
+// clang-format on
+
+#include "tensorflow/python/lib/core/ndarray_tensor_bridge.h"
 
 #include <vector>
 
 #include "tensorflow/c/c_api.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/python/lib/core/bfloat16.h"
-#include "tensorflow/python/lib/core/ndarray_tensor_bridge.h"
+#include "tensorflow/python/lib/core/py_util.h"
 
 namespace tensorflow {
 
@@ -214,6 +218,20 @@ Status ArrayFromMemory(int dim_size, npy_intp* dims, void* data, DataType dtype,
   }
   auto* np_array = reinterpret_cast<PyArrayObject*>(
       PyArray_SimpleNewFromData(dim_size, dims, type_num, data));
+  if (np_array == nullptr) {
+    string shape_str = absl::StrJoin(
+        absl::Span<npy_intp>{dims, static_cast<size_t>(dim_size)}, ", ");
+    if (PyErr_Occurred()) {
+      string exception_str = PyExceptionFetch();
+      PyErr_Clear();
+      return errors::InvalidArgument(
+          "Failed to create numpy array from tensor of shape [", shape_str,
+          "]. Numpy error: ", exception_str);
+    }
+    return errors::Internal(
+        "Failed to create numpy array from tensor of shape [", shape_str, "]");
+  }
+
   PyArray_CLEARFLAGS(np_array, NPY_ARRAY_OWNDATA);
   if (PyType_Ready(&TensorReleaserType) == -1) {
     return errors::Unknown("Python type initialization failed.");