# Function return type depends on an input flag : Bad?
**URL:** https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486
**Category:** Performance
**Tags:** zygote
**Created:** [June 24, 2021, 5:46am UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486 "2021-06-24T05:46:40Z")
**Posts on this page:** 16
**Page:** 1
### Author:  [@rakeshvar](https://discourse.julialang.org/u/rakeshvar)
#### Post date: [June 24, 2021, 5:46am UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/1 "2021-06-24T05:46:40Z")
I have a function that calculates ‘sky’ loss and optionally its gradient.
```julia
function _sky_loss_gradient(x, loss_only::Bool)
# calculation
loss_only && return loss
# more calculation
return loss, grad
end
sky_loss(x) = _sky_loss_gradient(x, true)
```
I was wondering if this is a good practice? As the compiler might not know the output type at run time. This is my use case: I am defining an autodiff reverse mode chain rule for backpropagation of loss gradients.
```julia
import ChainRulesCore: rrule, DoesNotExist, NO_FIELDS
function rrule(::typeof(sky_loss), x)
loss, grad = _sky_loss_gradient(exite, labels, false)
_pullback(x̄) = (NO_FIELDS, x̄.*grad, DoesNotExist())
loss, _pullback
end
```
Will the compiler be able to figure out that `sky_loss` returns a single value (not a tuple).
Or should I be using `Val(true)`? Something like…
```julia
function _sky_loss_gradient(x, ::Val{loss_only}) where loss_only
# calculation
loss_only && return loss
# more calculation
return loss, grad
end
sky_loss(x) = _sky_loss_gradient(x, Val(true))
```
Is there any other simpler way?
Edit: Users will never call `_sky_loss_gradient`. They will only use `sky_loss` and gradients are for autodiff via Zygote or Flux.
---
### Author:  [@Sukera](https://discourse.julialang.org/u/Sukera)
#### Post date: [June 24, 2021, 6:05am UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/2 "2021-06-24T06:05:39Z")
Yes, that is generally bad - it’s the definition of type instability (output type depending on the input _values_, not types) and moreover, users of your function won’t even know how many objects your function will return in the general case!
As your code is now, users will basically always have to do this dance:
```julia
res = sky_loss_gradient(inp, some_bool)
if res isa Tuple
# some code
# that
# may be very long
else
# ok we don't have a gradient..?
end
```
> [@rakeshvar](#):
>
> Will the compiler be able to figure out that `sky_loss` returns a single value (not a tuple).
If you only ever pass boolean literals into `sky_loss_gradient`, constant propagation should take care of this, yes. That being said, it’s probably best to handle those cases explicitly in `rrule` and seperate loss and gradient into two functions.
---
### Author:  [@rakeshvar](https://discourse.julialang.org/u/rakeshvar)
#### Post date: [June 24, 2021, 6:14am UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/3 "2021-06-24T06:14:26Z")
Thank you.
Users will never call `_sky_loss_gradient` (I added an underscore 🙂). The gradient is obtained only via Zygote or flux. So now can the compiler figure out that `sky_loss` does not return a tuple? `_sky_loss_gradient` is called explicitly only in two places once with a `true` and once with a `false`.
I can not write two functions as the calculations are highly intertwined and editing two functions parallely is buggy.
What do you think of `Val(true)` approach?
---
### Author:  [@jzr](https://discourse.julialang.org/u/jzr)
#### Post date: [June 24, 2021, 6:36am UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/4 "2021-06-24T06:36:16Z")
I think boolean flags are generally bad. I’d just have two functions. They don’t look intertwined to me, one can call the other and modify the result.
```julia
skyloss(x) = ...
function skyloss_and_grad(x)
loss = skyloss(x)
grad = ...
loss, grad
end
```
---
### Author:  [@gustaphe](https://discourse.julialang.org/u/gustaphe)
#### Post date: [June 24, 2021, 6:48am UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/5 "2021-06-24T06:48:50Z")
To add to this, if they _are_ intertwined, return both, and let the user do `(loss, ~) = ...` if they want to.
---
### Author:  [@rakeshvar](https://discourse.julialang.org/u/rakeshvar)
#### Post date: [June 24, 2021, 6:51am UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/6 "2021-06-24T06:51:27Z")
@gustaphe , Here `more calculation` depends on `calculation`. I am trying to avoid `more calculation`.
> [@rakeshvar](#):
>
> ```julia
> function _sky_loss_gradient(x, loss_only::Bool)
> # calculation
> loss_only && return loss
> # more calculation
> return loss, grad
> end
>
> ```
---
### Author:  [@Sukera](https://discourse.julialang.org/u/Sukera)
#### Post date: [June 24, 2021, 7:12am UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/7 "2021-06-24T07:12:39Z")
Let’s get at this from a different angle. In which situations will you need the gradient, and in which will you only need the loss? Maybe there’s some identifying characteristic that you can use dispatch with to get rid of the instability?
---
### Author:  [@aviatesk](https://discourse.julialang.org/u/aviatesk)
#### Post date: [June 24, 2021, 9:23am UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/8 "2021-06-24T09:23:20Z")
Well, as far as your “flag” is simple and can be resolved as compile-time constant, you can rely on constant prop’ performed by type inference:
```julia
julia> function foo(a, flag)
r = a > 0
if flag
return a, r
else
r
end
end
foo (generic function with 1 method)
julia> code_typed((Int,)) do n
foo(n, true) # compiler should fully understand the return type here
end |> first
CodeInfo(
1 ─ %1 = Base.slt_int(0, n)::Bool
│ %2 = Core.tuple(n, %1)::Tuple{Int64, Bool}
└── return %2
) => Tuple{Int64, Bool}
```
It’s not well documented which data can be constant-folded, but you can definitely use `Int`, `Bool`, `Symbol`, for example (N.B. the value of `String` won’t be constant-folded).
And IMHO there won’t be so many use cases of `Val` type, especially when it wraps a value that can be compile-time constant.
Having said that, if you can refactor your function so that its return types doesn’t depend on argument value, however it can be constant-folded, then it’s definitely better. Julia’s inference may give up constant prop’ when encountering a mutually recursive call cycle, for example, and debugging inference issues in such cases might be tricky.
---
### Author:  [@simeonschaub](https://discourse.julialang.org/u/simeonschaub)
#### Post date: [June 24, 2021, 10:13am UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/9 "2021-06-24T10:13:02Z")
> Having said that, if you can refactor your function so that its return types doesn’t depend on argument value, however it can be constant-folded, then it’s definitely better. Julia’s inference may give up constant prop’ when encountering a mutually recursive call cycle, for example, and debugging inference issues in such cases might be tricky.
Note that this is especially true when using AD. Zygote tends to generate code with a lot of closures, which really makes it difficult for the compiler to do constant propagation. That’s why I would generally recommend against relying on constant propagation too heavily when dealing with AD.
---
### Author:  [@pixel](https://discourse.julialang.org/u/pixel)
#### Post date: [June 24, 2021, 12:11pm UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/10 "2021-06-24T12:11:21Z")
Is this useful:
```julia
function _sky_loss_gradient(x, loss_only::Bool)
# calculation
loss_only && return loss,nothing
# more calculation
return loss, grad
end
```
the return type is always a tuple, and the boolean determines whether the second item in the tuple contains a useful value.
---
### Author:  [@Sukera](https://discourse.julialang.org/u/Sukera)
#### Post date: [June 24, 2021, 12:32pm UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/11 "2021-06-24T12:32:40Z")
It doesn’t change the fact that the type is unstable - you’ve just moved the instability from `typeof(loss)` vs. `Tuple{typeof(loss), typeof(grad)}` to `Tuple{typeof(loss), Nothing}` vs. `Tuple{typeof(loss), typeof(grad)}`. I guess the advantage is that there’s no question how many elements there are, so destructuring is an option now.
---
### Author:  [@Elrod](https://discourse.julialang.org/u/Elrod)
#### Post date: [June 24, 2021, 2:42pm UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/12 "2021-06-24T14:42:11Z")
Julia provides great tools for exploring type stability yourself:
```julia
function lossgrad(x, loss_only::Bool)
loss = 0.5sum(abs2, x)
loss_only && return loss
loss, x
end
function lossgrad(x, ::Val{loss_only}) where {loss_only}
loss = 0.5sum(abs2, x)
loss_only && return loss
loss, x
end
x = randn(10);
@code_warntype lossgrad(x, true)
@code_warntype lossgrad(x, Val(true))
@code_warntype lossgrad(x, Val(false))
```
---
### Author:  [@Sukera](https://discourse.julialang.org/u/Sukera)
#### Post date: [June 24, 2021, 3:43pm UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/13 "2021-06-24T15:43:43Z")
Yep - that only works if the `::Val{true}` is written in statically or derived from type information though, as far as I know. So having that `Val` be dynamic will (as far as I understand) still lead to type instability.
---
### Author:  [@Elrod](https://discourse.julialang.org/u/Elrod)
#### Post date: [June 24, 2021, 4:07pm UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/14 "2021-06-24T16:07:06Z")
Yes, but it avoids relying on const prop, which was the concern over:
```julia
function rrule(::typeof(sky_loss), x)
loss, grad = _sky_loss_gradient(exite, labels, false)
_pullback(x̄) = (NO_FIELDS, x̄.*grad, DoesNotExist())
loss, _pullback
end
```
---
### Author:  [@jzr](https://discourse.julialang.org/u/jzr)
#### Post date: [June 24, 2021, 5:52pm UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/15 "2021-06-24T17:52:44Z")
I think boolean flags and type instability are almost always bad design, regardless of their performance consequences. Make simple functions with clear interfaces. If you find that makes it faster too, that’s just gravy – the thing to get right first is design.
---
### Author:  [@rakeshvar](https://discourse.julialang.org/u/rakeshvar)
#### Post date: [July 7, 2021, 8:14am UTC](https://discourse.julialang.org/t/function-return-type-depends-on-an-input-flag-bad/63486/16 "2021-07-07T08:14:53Z")
Great. This is what I ended up doing.