fixes

Author: mcabbott

src/functions.jl

@@ -114,7 +114,7 @@ elseif VERSION >= v"1.1"
 end
 
 @static if VERSION > v"1.9-DEV"
-    # TODO: only supports dims=: for now
+    # TODO: support dims keyword. At present this will error with dims=1 etc.
     function Base.stack(A::KeyedArray; dims::Colon=:)
         data = @invoke stack(A::AbstractArray; dims)
         if !allequal(named_axiskeys(a) for a in A)

src/stack.jl

@@ -15,22 +15,22 @@ end
 
 # tuple of arrays
 function LazyStack.lazystack(x::Tuple{Vararg{<:KeyedArray}})
-    KeyedArray(LazyStack.stack(map(parent, x)), stack_keys(x))
+    KeyedArray(lazystack(map(parent, x)), stack_keys(x))
 end
 
 stack_keys(xs::Tuple{Vararg{<:KeyedArray}}) =
     (keys_or_axes(first(xs))..., Base.OneTo(length(xs)))
 
 # array of arrays: first strip off outer containers...
 function LazyStack.lazystack(xs::KeyedArray{<:AbstractArray})
-    KeyedArray(LazyStack.stack(parent(xs)), stack_keys(xs))
+    KeyedArray(lazystack(parent(xs)), stack_keys(xs))
 end
 function LazyStack.lazystack(xs::KeyedArray{<:AbstractArray,N,<:NamedDimsArray{L}}) where {L,N}
-    data = LazyStack.stack(parent(parent(xs)))
+    data = lazystack(parent(parent(xs)))
     KeyedArray(LazyStack.ensure_named(data, LazyStack.getnames(xs)), stack_keys(xs))
 end
 function LazyStack.lazystack(xs::NamedDimsArray{L,<:AbstractArray,N,<:KeyedArray}) where {L,N}
-    data = LazyStack.stack(parent(parent(xs)))
+    data = lazystack(parent(parent(xs)))
     LazyStack.ensure_named(KeyedArray(data, stack_keys(xs)), LazyStack.getnames(xs))
 end
 

test/runtests.jl

@@ -1,5 +1,5 @@
 using Test, AxisKeys, NamedDims
-using Statistics, OffsetArrays, Tables, UniqueVectors
+using Statistics, OffsetArrays, Tables, UniqueVectors, LazyStack
 using ChainRulesCore: ProjectTo, NoTangent
 using ChainRulesTestUtils: test_rrule
 using FiniteDifferences