@@ -2621,7 +2621,7 @@ function _typed_stack(::Colon, ::Type{T}, ::Type{S}, A, Aax=_axes(A)) where {T,
2621
2621
nothing === xit && return _empty_stack (:, T, S, A)
2622
2622
x1, _ = xit
2623
2623
ax1 = _axes (x1)
2624
- B = similar (_prototype (x1, A), T, ax1... , Aax... )
2624
+ B = similar (_first_array (x1, A), T, ax1... , Aax... )
2625
2625
off = firstindex (B)
2626
2626
len = length (x1)
2627
2627
while xit != = nothing
@@ -2665,7 +2665,7 @@ function _dim_stack(dims::Integer, ::Type{T}, ::Type{S}, A) where {T,S}
2665
2665
2666
2666
newaxis = _vec_axis (A)
2667
2667
outax = ntuple (d -> d== dims ? newaxis : _axes (x1)[d - (d> dims)], N1)
2668
- B = similar (_prototype (x1, A), T, outax... )
2668
+ B = similar (_first_array (x1, A), T, outax... )
2669
2669
2670
2670
iit = iterate (newaxis)
2671
2671
while xit != = nothing
@@ -2704,33 +2704,32 @@ end
2704
2704
end
2705
2705
2706
2706
# For `similar`, the goal is to stack an Array of CuArrays to a CuArray:
2707
- _prototype (x:: AbstractArray , A:: AbstractArray ) = x
2708
- _prototype (x:: AbstractArray , A) = x
2709
- _prototype (x, A:: AbstractArray ) = A
2710
- _prototype (x, A) = 1 : 0
2707
+ _first_array (x:: AbstractArray , ys... ) = x
2708
+ _first_array (x, ys... ) = _first_array (ys... )
2709
+ _first_array () = 1 : 0
2711
2710
2712
2711
# With tuple elements, we can make the empty array the right size:
2713
2712
function _empty_stack (:: Colon , :: Type{T} , :: Type{S} , A) where {T, S<: Tuple }
2714
- similar (_prototype ( nothing , A), T, OneTo (length (fieldtypes (S))), axes (A)... )
2713
+ similar (_first_array ( A), T, OneTo (length (fieldtypes (S))), axes (A)... )
2715
2714
end
2716
2715
function _empty_stack (dims:: Integer , :: Type{T} , :: Type{S} , A) where {T, S<: Tuple }
2717
2716
ax1 = OneTo (length (fieldtypes (S)))
2718
2717
dims in 1 : 2 || throw (ArgumentError (" cannot stack tuples along dims = $dims " ))
2719
- similar (_prototype ( nothing , A), T, ntuple (d -> d== dims ? OneTo (0 ) : ax1, 2 ))
2718
+ similar (_first_array ( A), T, ntuple (d -> d== dims ? OneTo (0 ) : ax1, 2 ))
2720
2719
end
2721
2720
# but with arrays of arrays, we must settle for the right ndims:
2722
2721
_empty_stack (dims, :: Type{T} , :: Type{S} , A) where {T,S} = _empty_stack (dims, T, IteratorSize (S), A)
2723
2722
_empty_stack (dims, :: Type{T} , :: HasLength , A) where {T} = _empty_stack (dims, T, HasShape {1} (), A)
2724
2723
_empty_stack (dims, :: Type{T} , :: IteratorSize , A) where {T} = _empty_stack (dims, T, HasShape {0} (), A)
2725
2724
2726
2725
function _empty_stack (:: Colon , :: Type{T} , :: HasShape{N} , A) where {T,N}
2727
- similar (_prototype ( nothing , A), T, ntuple (_-> OneTo (1 ), N)... , _axes (A)... )
2726
+ similar (_first_array ( A), T, ntuple (_-> OneTo (1 ), N)... , _axes (A)... )
2728
2727
end
2729
2728
function _empty_stack (dims:: Integer , :: Type{T} , :: HasShape{N} , A) where {T,N}
2730
2729
# Not sure we should check dims here, e.g. stack(Vector[]; dims=2) is an error
2731
2730
dims in 1 : N+ 1 || throw (ArgumentError (" cannot stack slices ndims(x) = $N along dims = $dims " ))
2732
2731
ax = ntuple (d -> d== dims ? _vec_axis (A) : OneTo (1 ), N+ 1 )
2733
- similar (_prototype ( nothing , A), T, ax... )
2732
+ similar (_first_array ( A), T, ax... )
2734
2733
end
2735
2734
2736
2735
# These make stack(()) work like stack([])
0 commit comments