@@ -262,6 +262,7 @@ impl<A, D> Array<A, D>
262
262
return Err ( ShapeError :: from_kind ( ErrorKind :: IncompatibleShape ) ) ;
263
263
}
264
264
265
+ let current_axis_len = self . len_of ( axis) ;
265
266
let remaining_shape = self . raw_dim ( ) . remove_axis ( axis) ;
266
267
let array_rem_shape = array. raw_dim ( ) . remove_axis ( axis) ;
267
268
@@ -281,22 +282,46 @@ impl<A, D> Array<A, D>
281
282
282
283
let self_is_empty = self . is_empty ( ) ;
283
284
284
- // array must be empty or have `axis` as the outermost (longest stride)
285
- // axis
286
- if !( self_is_empty ||
287
- self . axes ( ) . max_by_key ( |ax| ax. stride ) . map ( |ax| ax. axis ) == Some ( axis) )
288
- {
289
- return Err ( ShapeError :: from_kind ( ErrorKind :: IncompatibleLayout ) ) ;
285
+ // array must be empty or have `axis` as the outermost (longest stride) axis
286
+ if !self_is_empty && current_axis_len > 1 {
287
+ // `axis` must be max stride axis or equal to its stride
288
+ let max_stride_axis = self . axes ( ) . max_by_key ( |ax| ax. stride ) . unwrap ( ) ;
289
+ if max_stride_axis. axis != axis && max_stride_axis. stride > self . stride_of ( axis) {
290
+ return Err ( ShapeError :: from_kind ( ErrorKind :: IncompatibleLayout ) ) ;
291
+ }
290
292
}
291
293
292
294
// array must be be "full" (have no exterior holes)
293
295
if self . len ( ) != self . data . len ( ) {
294
296
return Err ( ShapeError :: from_kind ( ErrorKind :: IncompatibleLayout ) ) ;
295
297
}
298
+
296
299
let strides = if self_is_empty {
297
- // recompute strides - if the array was previously empty, it could have
298
- // zeros in strides.
299
- res_dim. default_strides ( )
300
+ // recompute strides - if the array was previously empty, it could have zeros in
301
+ // strides.
302
+ // The new order is based on c/f-contig but must have `axis` as outermost axis.
303
+ if axis == Axis ( self . ndim ( ) - 1 ) {
304
+ // prefer f-contig when appending to the last axis
305
+ // Axis n - 1 is outermost axis
306
+ res_dim. fortran_strides ( )
307
+ } else {
308
+ // Default with modification
309
+ res_dim. slice_mut ( ) . swap ( 0 , axis. index ( ) ) ;
310
+ let mut strides = res_dim. default_strides ( ) ;
311
+ res_dim. slice_mut ( ) . swap ( 0 , axis. index ( ) ) ;
312
+ strides. slice_mut ( ) . swap ( 0 , axis. index ( ) ) ;
313
+ strides
314
+ }
315
+ } else if current_axis_len == 1 {
316
+ // This is the outermost/longest stride axis; so we find the max across the other axes
317
+ let new_stride = self . axes ( ) . fold ( 1 , |acc, ax| {
318
+ if ax. axis == axis { acc } else {
319
+ Ord :: max ( acc, ax. len as isize * ax. stride )
320
+ }
321
+ } ) ;
322
+ let mut strides = self . strides . clone ( ) ;
323
+ strides[ axis. index ( ) ] = new_stride as usize ;
324
+ strides
300
325
} else {
301
326
self . strides . clone ( )
302
327
} ;
@@ -385,7 +410,8 @@ where
385
410
return ;
386
411
}
387
412
sort_axes_impl ( & mut a. dim , & mut a. strides , & mut b. dim , & mut b. strides ) ;
388
- debug_assert ! ( a. is_standard_layout( ) ) ;
413
+ debug_assert ! ( a. is_standard_layout( ) , "not std layout dim: {:?}, strides: {:?}" ,
414
+ a. shape( ) , a. strides( ) ) ;
389
415
}
390
416
391
417
fn sort_axes_impl < D > ( adim : & mut D , astrides : & mut D , bdim : & mut D , bstrides : & mut D )
0 commit comments