diff --git a/src/Dimensions/primitives.jl b/src/Dimensions/primitives.jl
index cf15dfc2e..9b224498c 100644
--- a/src/Dimensions/primitives.jl
+++ b/src/Dimensions/primitives.jl
@@ -346,7 +346,7 @@ function setdims end
 Swap dimensions for the passed in dimensions, in the
 order passed.
 
-Passing in the `Dimension` types rewraps the dimension index,
+Passing in the `Dimension` types rewraps the dimension lookup,
 keeping the index values and metadata, while constructed `Dimension`
 objects replace the original dimension. `nothing` leaves the original
 dimension as-is.
diff --git a/src/set.jl b/src/set.jl
index dac869a29..58df6acf7 100644
--- a/src/set.jl
+++ b/src/set.jl
@@ -114,27 +114,50 @@ julia> set(da, :custom => DD.Irregular(10, 12), Z => DD.Regular(9.9))
 function set end
 
 # Types are constructed
-Base.@assume_effects :effect_free set(x::DimArrayOrStack, ::Type{T}) where T = set(x, T())
-
+Base.@assume_effects :effect_free set(x::DimArrayOrStack, ::Type{T}) where T = 
+  set(x, T())
 # Dimensions and pairs are set for dimensions 
-Base.@assume_effects :effect_free set(A::DimArrayOrStack, args::Union{Dimension,DimTuple,Pair}...; kw...) =
-    rebuild(A, data(A), set(dims(A), args...; kw...))
+Base.@assume_effects :effect_free function set(
+  A::AbstractDimArray, args::Union{Dimension,DimTuple,Pair}...; kw...
+)
+    rebuild(A; dims=set(dims(A), args...; kw...))
+end
+Base.@assume_effects :effect_free function set(
+  st::AbstractDimStack, args::Union{Dimension,DimTuple,Pair}...; kw...
+)
+    ds = set(dims(st), args...; kw...)
+    if dimsmatch(ds, dims(st))
+        rebuild(st; dims=ds) 
+    else
+        dim_updates = map(rebuild, basedims(st), basedims(ds))
+        lds = map(layerdims(st)) do lds
+            # Swap out the dims with the updated dims
+            # that match the dims of this layer
+            map(val, dims(dim_updates, lds))
+        end
+        rebuild(st; dims=ds, layerdims=lds)
+    end
+end
 # Single traits are set for all dimensions
 Base.@assume_effects :effect_free set(A::DimArrayOrStack, x::LookupTrait) = 
     set(A, map(d -> basedims(d) => x, dims(A))...)
 # Single lookups are set for all dimensions
+# Need both for ambiguity
 Base.@assume_effects :effect_free set(A::AbstractDimArray, x::Lookup) = 
     set(A, map(d -> basedims(d) => x, dims(A))...)
 Base.@assume_effects :effect_free set(A::AbstractDimStack, x::Lookup) = 
     set(A, map(d -> basedims(d) => x, dims(A))...)
 # Arrays are set as data for AbstractDimArray
-Base.@assume_effects :effect_free set(A::AbstractDimArray, newdata::AbstractArray) = begin
+Base.@assume_effects :effect_free function set(
+    A::AbstractDimArray, newdata::AbstractArray
+)
     axes(A) == axes(newdata) || _axiserr(A, newdata)
     rebuild(A; data=newdata)
 end
-
 # NamedTuples are set as data for AbstractDimStack
-Base.@assume_effects :effect_free set(s::AbstractDimStack, newdata::NamedTuple) = begin
+Base.@assume_effects :effect_free function set(
+    s::AbstractDimStack, newdata::NamedTuple
+)
     dat = data(s)
     keys(dat) === keys(newdata) || _keyerr(keys(dat), keys(newdata))
     map(dat, newdata) do d, nd
diff --git a/test/set.jl b/test/set.jl
index 2b87ed939..e4ce8e999 100644
--- a/test/set.jl
+++ b/test/set.jl
@@ -2,6 +2,7 @@ using DimensionalData, Test
 using DimensionalData.Lookups, DimensionalData.Dimensions
 
 using DimensionalData.Lookups: _set
+using DimensionalData: layerdims
 
 a = [1 2; 3 4]
 dimz = (X(143.0:2.0:145.0; lookup=Sampled(order=ForwardOrdered()), metadata=Metadata(Dict(:meta => "X"))),
@@ -32,13 +33,17 @@ end
 end
 
 @testset "DimStack Dimension" begin
-    @test typeof(dims(set(s, row=X, column=Z))) <: Tuple{<:X,<:Z}
+    s1 = set(s, row=X, column=Z)
+    @test typeof(dims(s1)) <: Tuple{<:X,<:Z}
+    @test layerdims(s1) == (; test2=(X(), Z()), test3=(X(), Z()))
     @test typeof(dims(set(s, row=X(), column=Z()))) <: Tuple{<:X,<:Z}
-    @test typeof(dims(set(s, row=:row2, column=:column2))) <: Tuple{<:Dim{:row2},<:Dim{:column2}}
+    s1 = set(s, row=:row2, column=:column2)
+    @test typeof(dims(s1)) <: Tuple{<:Dim{:row2},<:Dim{:column2}}
+    @test layerdims(s1) == (; test2=(Dim{:row2}(), Dim{:column2}()), test3=(Dim{:row2}(), Dim{:column2}()))
     @test typeof(dims(set(s, :column => Ti(), :row => Z))) <: Tuple{<:Z,<:Ti}
     @test typeof(dims(set(s, Dim{:row}(Y()), Dim{:column}(X())))) <: Tuple{<:Y,<:X}
     @test typeof(dims(set(s, (Dim{:row}(Y), Dim{:column}(X))))) <: Tuple{<:Y,<:X}
-    @test index(set(s, Dim{:row}([:x, :y, :z])), :row) == [:x, :y, :z] 
+    @test parent(lookup(set(s, Dim{:row}([:x, :y, :z])), :row)) == [:x, :y, :z] 
 end
 
 @testset "DimArray Dimension" begin