How to enforce minimum dimensions for arrays that might be scalars?

I have a function expecting a Matrix and a Vector argument:

function f(m::Matrix, v::Vector) 
...
end

Inside f performs some linear algebra. My problem: I need f to work when its arguments are scalars as well. How should I handle this?

If this were numpy I might wrap the arguments in numpy arrays and use the ndim= keyword argument to enforce the shape of each. Does Julia have a similar facility?

If the operations ... will work equally well for numbers, then you can define the function as e.g. f(m::AbstractMatrix, v::Union{Number, AbstractVector}) or even just f(m, v). A different version will still be compiled for each combination of concrete types.

If not, and you need a separate implementation, and it’s cleanest to make another method f(m::Number, v::Number) to contain this.

You can re-use the matrix implementation by writing perhaps f(m::Number, v::Number) = f(hcat(m), [v]) |> only i.e. creating a 1x1 matrix and a vector, calling the first method, and then (perhaps) unwrapping back to a number. But very likely you can do something more efficient working directly with scalars instead.

9 Likes

Thanks!

It’s very common in Matlab and NumPy to write functions which “vectorize” over one or more arguments, supporting scalar or vector operations, where the vector version just involves applying the scalar version to each element. This is rarely necessary or helpful in Julia. If that’s what you’re doing here (it’s hard to tell from your description), then I would suggest avoiding the issue completely and only implementing f(x::AbstractMatrix, v::Number) and then using broadcasting when you call your function to apply it elementwise.

For example, I’m imagining your function is something like:

julia> function f_scalar(x, v)
         v * x
       end
f_scalar (generic function with 1 method)

julia> function f_vector(x, v)
         [v_i * x for v_i in v]
       end
f_vector (generic function with 1 method)

where f_vector just applies the same operation for each element in the vector v.

If this is what you’re trying to achieve, then I would suggest that you not write f_vector at all. All you need is f_scalar, and you can apply it to v of any dimension via broadcasting with .

# Scalar `v`, returns a single `Matrix` result
julia> f_scalar(ones(2, 2), 1.5)
2×2 Array{Float64,2}:
 1.5  1.5
 1.5  1.5

# Vector `v`. Returns a vector of matrices, one for each element of `v`
julia> f_scalar.(Ref(ones(2, 2)), [1.5, 2.5])
2-element Array{Array{Float64,2},1}:
 [1.5 1.5; 1.5 1.5]
 [2.5 2.5; 2.5 2.5]

# Matrix `v`. Returns a matrix of matrices, one for each element of `v`:
julia> f_scalar.(Ref(ones(2, 2)), [1.5 2.5; 3.5 4.5])
2×2 Array{Array{Float64,2},2}:
 [1.5 1.5; 1.5 1.5]  [2.5 2.5; 2.5 2.5]
 [3.5 3.5; 3.5 3.5]  [4.5 4.5; 4.5 4.5]

The Ref means to treat the first argument as a scalar, passing it whole to each call to f_scalar.

This is a much more common approach in Julia, and it results in more efficient code with less work. For more info, see: More Dots: Syntactic Loop Fusion in Julia

7 Likes