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