3

How do I use both Julia's dot notation to do elementwise operations AND ensure that the result is saved in an already existing array?

function myfun(x, y)
    return x + y
end

a = myfun(1, 2)  # Results in a == 3

a = myfun.([1 2], [3; 4])  # Results in a == [4 5; 5 6]

function myfun!(x, y, out)
    out .= x + y
end

a = zeros(2, 2)
myfun!.([1 2], [3; 4], a)  # Results in a DimensionMismatch error

Also, does @. a = myfun([1 2], [3; 4]) write the result to a in the same way as I am trying to achieve with myfun!()? That is, does that line write the result directly to a without saving storing the result anywhere else first?

Fredrik P
  • 682
  • 1
  • 8
  • 21
  • 1
    `out` should preferably be the first parameter, see https://docs.julialang.org/en/v1/manual/style-guide/#Write-functions-with-argument-ordering-similar-to-Julia-Base. – phipsgabler Feb 22 '22 at 11:28
  • 1
    If your function is 'inherently scalar', i.e. is most naturally expressed as a scalar function (as in your example), then the most idiomatic approach is to _not_ define `myfun!`, but instead use the `a .= myfun.(x, y)` alternative. Making a `myfun!` is what you would do if it's difficult or awkward to achieve in-place operations otherwise, which is not the case here. Then you push the broadcasting to a higher level, rather then a lower. – DNF Feb 22 '22 at 15:48

2 Answers2

4

Your code should be:

julia> function myfun!(x, y, out)
           out .= x .+ y
       end
myfun! (generic function with 1 method)

julia> myfun!([1 2], [3; 4], a)
2×2 Matrix{Float64}:
 4.0  5.0
 5.0  6.0

julia> a
2×2 Matrix{Float64}:
 4.0  5.0
 5.0  6.0

As for @. a = myfun([1 2], [3; 4]) - the answer is yes, it does not create temporary arrays and operates in-place.

Bogumił Kamiński
  • 66,844
  • 3
  • 80
  • 107
  • I infer that it is not possible to do what I am trying to: that is, to treat the `out` argument in the same element-wise fashion as the `x` and `y` arguments when doing in-place operations. Is this correct? – Fredrik P Feb 22 '22 at 10:17
  • 1
    When you write `myfun!.([1 2], [3; 4], a)` you do not pass `a` to `myfun!` but only its elements. Does this clarify your issue? – Bogumił Kamiński Feb 22 '22 at 10:29
  • Yes, it does :-) – Fredrik P Feb 22 '22 at 13:37
2

This isn't commonly required, and there are usually better ways to achieve this, but it's possible to broadcast on an output argument by using Reference values that point inside the output array.

julia> a = zeros(2, 2)
2×2 Matrix{Float64}:
 0.0  0.0
 0.0  0.0

julia> function myfun!(out, x, y)
          out[] = x + y
       end
myfun! (generic function with 1 method)

julia> myfun!.((Ref(a, i) for i in LinearIndices(a)), [1 2], [3; 4])
2×2 Matrix{Int64}:
 4  5
 5  6

julia> a
2×2 Matrix{Float64}:
 4.0  5.0
 5.0  6.0

Edit: Changed out to be the first parameter as per the Style guide - thanks to @phipsgabler for the reminder.

Sundar R
  • 13,776
  • 6
  • 49
  • 76