97
97
res_dim. set_axis ( axis, 0 ) ;
98
98
let mut res;
99
99
unsafe {
100
+ // create the axis so that 0 is outermost axis
100
101
res_dim. slice_mut ( ) . swap ( axis. index ( ) , 0 ) ;
101
102
res = Array :: from_shape_vec_unchecked ( res_dim, Vec :: with_capacity ( new_len) ) ;
102
- res. swap_axes ( axis. index ( ) , 0 ) ;
103
+ // swap axes so that `axis` is outermost axis
104
+ res. swap_axes ( 0 , axis. index ( ) ) ;
103
105
}
104
106
105
107
for array in arrays {
@@ -139,7 +141,7 @@ pub fn stack_new_axis<A, D>(
139
141
arrays : & [ ArrayView < A , D > ] ,
140
142
) -> Result < Array < A , D :: Larger > , ShapeError >
141
143
where
142
- A : Copy ,
144
+ A : Clone ,
143
145
D : Dimension ,
144
146
D :: Larger : RemoveAxis ,
145
147
{
@@ -159,24 +161,25 @@ where
159
161
160
162
res_dim. set_axis ( axis, arrays. len ( ) ) ;
161
163
162
- // we can safely use uninitialized values here because we will
163
- // overwrite every one of them.
164
- let mut res = Array :: uninit ( res_dim) ;
165
-
166
- res. axis_iter_mut ( axis)
167
- . zip ( arrays. iter ( ) )
168
- . for_each ( |( assign_view, array) | {
169
- // assign_view is D::Larger::Smaller which is usually == D
170
- // (but if D is Ix6, we have IxD != Ix6 here; differing types
171
- // but same number of axes).
172
- let assign_view = assign_view. into_dimensionality :: < D > ( )
173
- . expect ( "same-dimensionality cast" ) ;
174
- array. assign_to ( assign_view) ;
175
- } ) ;
164
+ let new_len = dimension:: size_of_shape_checked ( & res_dim) ?;
176
165
166
+ // start with empty array with precomputed capacity
167
+ res_dim. set_axis ( axis, 0 ) ;
168
+ let mut res;
177
169
unsafe {
178
- Ok ( res. assume_init ( ) )
170
+ // create the axis so that 0 is outermost axis
171
+ res_dim. slice_mut ( ) . swap ( axis. index ( ) , 0 ) ;
172
+ res = Array :: from_shape_vec_unchecked ( res_dim, Vec :: with_capacity ( new_len) ) ;
173
+ // swap axes so that `axis` is outermost axis
174
+ res. swap_axes ( 0 , axis. index ( ) ) ;
179
175
}
176
+
177
+ for array in arrays {
178
+ res. try_append_array ( axis, array. clone ( ) . insert_axis ( axis) ) ?;
179
+ }
180
+
181
+ debug_assert_eq ! ( res. len_of( axis) , arrays. len( ) ) ;
182
+ Ok ( res)
180
183
}
181
184
182
185
/// Stack arrays along the new axis.
0 commit comments