# How can I make an in-place swap function?

Suppose I have two arrays `A, B` of different sizes. Then `A .+ B` broadcasts them to a common size (if possible) and sums them. Suppose `A` here is “larger” (in the sense that `all(size(B) .== size(A) .|| size(B) .== 1)` is true).

Given an index `i` of `A` (`A[i]` ), how can I compute `j` such that `B[j]` gives the element of `B` that would be summed to `A[i]` in the broadcasted expression?

My motivation for this question, is that I am trying to code an in-place version of the following function:

``````function swap(conditions::Union{AbstractArray{Bool}, Bool}, x, y)
x_new = ifelse.(conditions, y, x)
y_new = ifelse.(conditions, x, y)
return x_new, y_new
end
``````

The following:

``````function swap!(x::AbstractArray, y::AbstractArray, conditions::AbstractArray{Bool})
if size(x) == size(y) == size(conditions)
for i in eachindex(x, y, conditions)
if conditions[i]
x[i], y[i] = y[i], x[i]
else
x[i], y[i] = x[i], y[i]
end
end
return x, y
else
throw(DimensionMismatch("x, y and conditions must have the same size"))
end
end
``````

works if x, y, conditions are of the same size. I would like to relax that, requiring that `x,y` be of the same size, but `conditions` need only be “broadcastable” to a matching size.

Hm good question, the closest I could come up with quickly was this, but it still allocates an array of `nothing`s which is not nice

``````function swap!(x, y, conditions)
if c
x[i], y[i] = y[i], x[i]
end
end
return
end
``````

Maybe a `FillArray` could help, though I’m not sure how.

I found some broadcasting internals that I could coax into doing this. Maybe there’s a different more obvious way

``````julia> function swap!(x, y, conditions)
for (i, c) in b
if c
x[i], y[i] = y[i], x[i]
end
end
return
end
swap! (generic function with 2 methods)

julia> x = fill(1, 3, 5); y = fill(2, 3, 5); conditions = [true, false, true, false, true]';

julia> swap!(x, y, conditions)

julia> x
3×5 Matrix{Int64}:
2  1  2  1  2
2  1  2  1  2
2  1  2  1  2

julia> y
3×5 Matrix{Int64}:
1  2  1  2  1
1  2  1  2  1
1  2  1  2  1

julia> @time swap!(x, y, conditions)
0.000012 seconds
``````
4 Likes

Here’s what I came up with.The first uses a `view` into `conditions` that repeatedly samples the dimensions that need broadcasting. The second figures out which dimensions need to be broadcast, then adjusts the index into `conditions` accordingly.

``````function swapif1!(x,y,cond)
axes(x) == axes(y) || error("x and y must have the same axes")
all(d -> size(cond,d) == 1 || axes(cond,d) == axes(x,d), ndims(x)) || error("could not broadcast cond to size of x, y")
axc = ntuple(d -> ifelse(size(cond,d) == 1, StepRangeLen(firstindex(cond,d),0,size(x,d)), :), Val(ndims(x)))
vcond = view(cond,axc...)
for i in eachindex(x,y,vcond)
if vcond[i]
x[i],y[i] = y[i],x[i]
end
end
return x,y
end

function swapif2!(x,y,cond)
axes(x) == axes(y) || error("x and y must have the same axes")
all(d -> size(cond,d) == 1 || axes(cond,d) == axes(x,d), ndims(x)) || error("could not broadcast cond to size of x, y")
szc1 = ntuple(d -> size(cond,d) == 1, Val(ndims(x))) # find axes where size(cond,ax) == 1
indc1 = ntuple(d -> firstindex(cond,d), Val(ndims(x))) # valid index for c on axes of size 1
for i in eachindex(IndexCartesian(), x, y)
ic = CartesianIndex(ifelse.(szc1, indc1, Tuple(i))...)
if cond[ic]
x[i],y[i] = y[i],x[i]
end
end
return x,y
end
``````

The first one ends up allocating a lot. I didn’t look into why, but there’s probably a type instability somewhere. So I’ll recommend the second.

But I think I like the suggestion above that uses tools from `Broadcast` to handle this. That seems clean.

1 Like

Very elegant, thanks!

This is not public API, though, I think?

Not sure. Not everything considered public/stable needs to be exported I think.

It looks like `Broadcasted` appears in the docs here. While it is not documented formally and explicitly, it is documented as part of the broadcasting interface. It is intended to be used to extend the broadcast machinery, so (in my reading) it is part of the broadcasting API and thus part of the Julia API, even though it is not exported.

However, I would use

``````Broadcast.broadcasted(tuple, CartesianIndices(x), conditions)
``````

rather than

``````Broadcast.Broadcasted(tuple, (CartesianIndices(x), conditions))
``````

The former produces the latter, but appears to be the more standard interface. Notice

``````julia> @code_lowered tuple.(CartesianIndices(x), conditions)
CodeInfo(
1 ─ %1 = Base.broadcasted(Main.tuple, x1, x2)
│   %2 = Base.materialize(%1)
└──      return %2
)
``````

So what is being done here is equivalent to the `broadcasted` version, except that a normal call with dots also calls `materialize` to collect the result to a variable. But the non-`materialize`d version can still be iterated and that’s why it works in this `for` loop. `Broadcast.broadcasted` can also be useful if you want to assemble a broadcasted expression in multiple parts or over multiple lines without materializing the intermediate results.

2 Likes

Good point. So I’m writing it like this now:

``````function swap!(conditions::Union{AbstractArray{Bool}, Bool}, x::AbstractArray, y::AbstractArray)
if size(x) == size(y)
b = broadcasted(tuple, CartesianIndices(x), CartesianIndices(y), conditions)
for (ix, iy, c) in b
if c
x[ix], y[iy] = y[ix], x[iy]
end
end
return x, y
else
throw(DimensionMismatch("x, y must have the same size"))
end
end
``````

I am using different indices `ix, iy` in case `x, y` have different indexing patterns (though not sure this is the best approach).

You could just use ordinary broacasting via StructArrays.jl:

``````using StructArrays
swap(condition::Bool, (x, y)) = condition ? (y, x) : (x, y)
function swap!(x::AbstractArray, y::AbstractArray, conditions)
a = StructArray((x, y))
a .= swap.(conditions, a)
return x, y
end
``````
3 Likes