Specifying output container type based on input container

General Usage
1

I have a question about generic functions which is probably something simple I am missing. Let’s say I have a function f defined as follows:

function f(x::A)::AbstractArray{T, N-1} where {T<:Real, N, A<:AbstractArray{T, N}}
    dropdims(sum(x, dims=(1,)), dims=(1,))
end

This function takes an arbitrary array x as input, and specifies to the compiler that an array will be returned ::AbstractArray{T,N-1}, i.e., with the same parameters T, and 1 fewer dimension.

However, is there a way I can specify that an array of the same container type as x will always be returned? For example, if x is an Array{T,2}, I want to indicate that an Array{T,1} will also be returned - not just any AbstractArray. Similar for StaticArrays, etc.

I tried the obviously incorrect expression

function f(x::A)::A{T,N-1} where {T<:Real, N, A<:AbstractArray{T, N}}
    dropdims(sum(x, dims=(1,)), dims=(1,))
end

which does not work. Is there a way to do this?

I should also specify - for this simple function I am sure the compiler can figure it out. But for more complex functions in my codebase I find it helped to indicate the output type… so perhaps indicating the correct container type will help as well.

Thanks!
Miles

2

Here’s a hacky attempt to inform the compiler about the container type (using How to get the container type of a container?):

constructor_of(::Type{T}) where T =
    getfield(T.name.module, Symbol(T.name.name))

function f(x::A) where {T<:Real, N, A<:AbstractArray{T, N}}
   dropdims(sum(x, dims=(1,)), dims=(1,))::constructor_of(A){T,N-1}
end

where you would wrap the entire function block with ::constructor_of(A){T,N-1}.

But surely there’s a way to do it directly in the function signature?

One other attempt:

function f(x::A)::(A.name.wrapper){T,N-1} where {T<:Real, N, A<:AbstractArray{T, N}}
   dropdims(sum(x, dims=(1,)), dims=(1,))
end

this feels illegal though. Not sure if A.name.wrapper is stable or not?

3

No. There is no API for specifying what the element type of an output array (or its dimensionality) ought to be because that’s not something inherent to the type system you can specify for dispatch (which doesn’t include the return type either way - that’s only part of what is inferred as an optimization).

4

There is no API for specifying what the element type of an output array (or its dimensionality) ought to be because that’s not something inherent to the type system you can specify for dispatch

Sorry, not sure I follow. Are you stating this as a terminology thing, or are you actually saying the effect of :: has no effect on the function? Because, for example:

f(x::AbstractArray{Float64})::AbstractArray{Float32} = x

f(randn(Float64, 100)) # this is now of type Array{Float32}

Regardless of this finer point, the reason I want to specify container type is simply to help out the compiler for functions where the output type is ambiguous, and it might not be able to infer it by itself.

5

that syntax is doing conversion

julia> Meta.@lower f(x::AbstractArray{Float64})::AbstractArray{Float32} = x
:($(Expr(:thunk, CodeInfo(
    @ none within `top-level scope`
1 ─      $(Expr(:thunk, CodeInfo(
    @ none within `top-level scope`
1 ─     return $(Expr(:method, :f))
)))
│        $(Expr(:method, :f))
│   %3 = Core.Typeof(f)
│   %4 = Core.apply_type(AbstractArray, Float64)
│   %5 = Core.svec(%3, %4)
│   %6 = Core.svec()
│   %7 = Core.svec(%5, %6, $(QuoteNode(:(#= REPL[1]:1 =#))))
│        $(Expr(:method, :f, :(%7), CodeInfo(
    @ REPL[1]:1 within `none`
1 ─ %1 = Core.apply_type(AbstractArray, Float32)
│   %2 = Base.convert(%1, x)
│   %3 = Core.typeassert(%2, %1)
└──      return %3
6

Yes, I know this, sorry if I wasn’t clear. Of course I could always manually wrap convert((A.name.wrapper){T,N-1}, ...) around my function block, it’s just ugly. And the other thing about the f(...)::T syntax I like is it acts as a form of documentation to other users.

I am basically looking for a cleaner way to specify it in the signature without a hacky solution like f(x::A)::(A.name.wrapper){T,N-1} where .... If nothing else exists I guess I can go with that.

7

because convert is not free, it may not be a good idea because of performance, I don’t know if you want correctness or performance or simply an API guarantee you think would be nice to have.

In the correctness + performance case, when you have total control, I do something like this:

this is doing a typeassert, which makes sure I get what I expected and conversion never happens (when it’s transferred to a typed buffer container which needs type prediction ahead of time). But if you don’t have total control, it’s hard to say.

I would think CUDA.jl-world has similar problem and they may have a solution utility

1 Like
8

Thanks. I guess there is no built-in syntax to do it at the function signature level? Some of the functions I am dealing with have multiple return statements so it could get ugly quickly

Btw, all of my functions already satisfy the type I am expecting; it is just sometimes the compiler can’t figure it out in the locations the function is being called. i.e., the ::(A.name.wrapper){T,N-1} syntax should not impact performance as it wouldn’t be doing conversions, it would just tell the compiler the return type. (But there’s no cleaner way to write this, I guess)

julia> @code_llvm convert(Array{Float64}, randn(Float64, 100))
;  @ array.jl:617 within `convert`
define nonnull {}* @julia_convert_1423({}* readonly %0, {}* nonnull align 16 dereferenceable(40) %1) #0 {
top:
  ret {}* %1
}
1 Like
9

right, you need to attached ::T not at the function definition level because that’s a convert, you need it at

return res::T

to make it a type assert, and “get the concrete container type of the <:AbstractArray” is not a well defined thing in the type system unfortunately

1 Like
10

They are identical if the type is correct though, right?

e.g., signature level:

julia -e 'function g(x::A)::(A.name.wrapper{T,N-1}) where {T,N,A<:AbstractArray{T,N}}; dropdims(x, dims=(1,)); end; x = ones(32, 1); import InteractiveUtils: @code_llvm; @code_llvm g(x)' \
    2>&1 > signature_level.txt

Function level:

julia -e 'function g(x::A) where {T,N,A<:AbstractArray{T,N}}; dropdims(x, dims=(1,))::(A.name.wrapper{T,N-1}); end; x = ones(32, 1); import InteractiveUtils: @code_llvm; @code_llvm g(x)' \
    2>&1 > function_level.txt

The diff is just different addresses:

<      store atomic i64 140599633651680, i64* %15 unordered, align 8
---
>      store atomic i64 139832940386272, i64* %15 unordered, align 8
45c45
<      store {}* inttoptr (i64 140599707586240 to {}*), {}** %16, align 8
---
>      store {}* inttoptr (i64 139833014320832 to {}*), {}** %16, align 8
93c93
<       store atomic i64 140599633764176, i64* %28 unordered, align 8
---
>       store atomic i64 139832940498768, i64* %28 unordered, align 8
97c97
<       %30 = call nonnull {}* inttoptr (i64 140600165136976 to {}* ({}*, {}*, {}*)*)({}* inttoptr (i64 140599635100832 to {}*), {}* nonnull %0, {}* nonnull %26)
---
>       %30 = call nonnull {}* inttoptr (i64 139833471871568 to {}* ({}*, {}*, {}*)*)({}* inttoptr (i64 139832941835424 to {}*), {}* nonnull %0, {}* nonnull %26)

I think it’s nicer when you have multiple return statements, and also for documentation if a user is looking at the signature, maybe. I wish there was cleaner syntax though…

11

typeassert would error if it doesn’t work, convert may silently cost you performance. again, I guess it depends on your goal

1 Like
12

I see, that makes sense to me. Thanks!

13

This is a more refined version (from ConstructionBase.jl) which completely avoids Julia internals:

@generated function constructorof(::Type{T}) where T
    getfield(parentmodule(T), nameof(T))
end

I think ultimately the issue is that, for diagonal types like f(x::A)::A where A, when f is called the type variable A will be concrete and fully parameterized; before you can change some parameters, you first need to shed some by climbing up to a non-concrete UnionAll type.

Without this type annotation/conversion, f’s output type is being correctly inferred on my machine. I don’t think there’s any need for concern one way or the other, but are you sure this is the point in your codebase where type stability is being lost?

1 Like
14

Thanks!

Without this type annotation/conversion, f’s output type is being correctly inferred on my machine. I don’t think there’s any need for concern one way or the other, but are you sure this is the point in your codebase where type stability is being lost?

Yeah. The tricky thing with my codebase SymbolicRegression.jl is that it is set up to allow arbitrary user-defined operators, loss functions, and constraints. However, sometimes user-defined operators introduce their own type stability for whatever reason (a lot of new Julia users are introduced via the PySR frontend, wanting to customize things for their problem), so I have noticed that it helps a lot if I just enforce type stability on their behalf. Then, the user won’t experience a weird slow-down and have to learn these advanced concepts in Julia just to use my library efficiently.

Right now SymbolicRegression.jl only allows for scalar operators, but I’m slowly expanding type constraints to allow for vector operators. So that’s where my question about specifying container types comes from.

15

Maybe another strategy is to use Core.Compiler.return_type or Test.@inferred to warn or disallow the user from introducing these type instabilities at all. Would lead them to correcting internal instabilities which presumably would make the code even faster than just annotating the return type for them, although would be more work on their end, so certainly a tradeoff.

1 Like