@@ -459,7 +459,8 @@ PyArray_ConcatenateFlattenedArrays(int narrays, PyArrayObject **arrays,
459
459
PyTypeObject * subtype = & PyArray_Type ;
460
460
double priority = NPY_PRIORITY ;
461
461
int iarrays ;
462
- npy_intp shape [2 ], strides [2 ];
462
+ npy_intp stride , sizes [NPY_MAXDIMS ];
463
+ npy_intp shape = 0 ;
463
464
PyArray_Descr * dtype = NULL ;
464
465
PyArrayObject * ret = NULL ;
465
466
PyArrayObject_fields * sliding_view = NULL ;
@@ -470,19 +471,17 @@ PyArray_ConcatenateFlattenedArrays(int narrays, PyArrayObject **arrays,
470
471
return NULL ;
471
472
}
472
473
473
- /* All the arrays must have the same total number of elements */
474
- shape [0 ] = narrays ;
475
- shape [1 ] = PyArray_SIZE (arrays [0 ]);
476
-
477
474
/*
478
475
* Figure out the final concatenated shape starting from the first
479
476
* array's shape.
480
477
*/
481
- for (iarrays = 1 ; iarrays < narrays ; ++ iarrays ) {
482
- if (PyArray_SIZE (arrays [iarrays ]) != shape [1 ]) {
478
+ for (iarrays = 0 ; iarrays < narrays ; ++ iarrays ) {
479
+ shape += sizes [iarrays ] = PyArray_SIZE (arrays [iarrays ]);
480
+ /* Check for overflow */
481
+ if (shape < 0 ) {
483
482
PyErr_SetString (PyExc_ValueError ,
484
- "all the input arrays must have same "
485
- "number of elements " );
483
+ "total number of elements "
484
+ "too large to concatenate " );
486
485
return NULL ;
487
486
}
488
487
}
@@ -504,15 +503,14 @@ PyArray_ConcatenateFlattenedArrays(int narrays, PyArrayObject **arrays,
504
503
return NULL ;
505
504
}
506
505
507
- strides [1 ] = dtype -> elsize ;
508
- strides [0 ] = strides [1 ] * shape [1 ];
506
+ stride = dtype -> elsize ;
509
507
510
508
/* Allocate the array for the result. This steals the 'dtype' reference. */
511
509
ret = (PyArrayObject * )PyArray_NewFromDescr (subtype ,
512
510
dtype ,
513
- 2 ,
514
- shape ,
515
- strides ,
511
+ 1 ,
512
+ & shape ,
513
+ & stride ,
516
514
NULL ,
517
515
0 ,
518
516
NULL );
@@ -530,9 +528,10 @@ PyArray_ConcatenateFlattenedArrays(int narrays, PyArrayObject **arrays,
530
528
Py_DECREF (ret );
531
529
return NULL ;
532
530
}
533
- /* Each array gets flattened into one slot along 'axis' */
534
- sliding_view -> dimensions [0 ] = 1 ;
531
+
535
532
for (iarrays = 0 ; iarrays < narrays ; ++ iarrays ) {
533
+ /* Adjust the window dimensions for this array */
534
+ sliding_view -> dimensions [0 ] = sizes [iarrays ];
536
535
537
536
/* Copy the data for this array */
538
537
if (PyArray_CopyAsFlat ((PyArrayObject * )sliding_view , arrays [iarrays ],
@@ -543,7 +542,7 @@ PyArray_ConcatenateFlattenedArrays(int narrays, PyArrayObject **arrays,
543
542
}
544
543
545
544
/* Slide to the start of the next window */
546
- sliding_view -> data += sliding_view -> strides [0 ];
545
+ sliding_view -> data += sliding_view -> strides [0 ] * sizes [ iarrays ] ;
547
546
}
548
547
549
548
Py_DECREF (sliding_view );
0 commit comments