Assign values to `slicedim`


Is there anyway to assign values to a section of an array?

For example

using Distributions, StatsFuns

# generate an array
y = rand(Normal(), 4, 3, 2)

# section of y to modified
modify = slicedim(y, 3, 2)

# modify section
modified = softplus.(modify)

# would like to do something like this
slicedim(y, 3, 2) = modified

Something like this is obviously possible

y[:, :, 2] = softplus.(y[:, :, 2])

But I like the syntax of the slicedim approach to be more “programmatic”


The code for slicedim() is just (from base/abstractarraymath.jl):

function slicedim(A::AbstractArray, d::Integer, i)
    d >= 1 || throw(ArgumentError("dimension must be ≥ 1"))
    nd = ndims(A)
    d > nd && (i == 1 || throw_boundserror(A, (ntuple(k->Colon(),nd)..., ntuple(k->1,d-1-nd)..., i)))
    A[setindex(indices(A), i, d)...]

You could write your own function slicedimview which runs exactly the same code but returns a view into the resulting array. Then you could do:

slicedimview(y, 3, 2) .= modified

However: Note that by splitting up the code into multiple statements, you’re forcing Julia to create a new temporary array to hold modified. If you do the entire computation as a single expression, then you’ll benefit from dot fusion and be able to avoid allocating that temporary array. That is, you can do:

y[:, :, 2] .= softplus.(@view y[:, :, 2])

which will not allocate any temporary arrays. Or, with the hypothetical slicedimview() function, you could do:

slicedimview(y, 3, 2) .= softplus.(slicedimview(y, 3, 2))

and also not allocate any temporary arrays.


Thanks! This is what I was looking for.