Add a 2-arg flavor, argmax(f, domain)

[drewrobson]

With the rename to argmax, and dims now a kwarg, it seems very tempting to me to define a 2-arg flavor, argmax(f, itr), where itr is an iterable subset of the domain of f. This would return an x in itr that maximizes f(x), which is quite literally the mathematical definition of argmax.

(originally proposed here)

Edit: Previous wording may have been confusing, so I have attempted to improve it.

[drewrobson]

If we have #27612, then this is as simple as defining the following:

Base.argmax(f, itr) = Base.argmax(f(x) for x in itr)

(but I haven't checked what to do about dims yet)

[ararslan]

A more correct definition would be to adapt the current definitions of findmin and findmax to accept a function argument, then simply adjust the current argmax definition to

Base.argmax(f, x::AbstractArray; dims=:) = findmax(f, x, dims=dims)[2]
Base.argmax(f, itr) = findmax(f, itr)[2]

with the addition of

Base.argmax(x::AbstractArray; dims=:) = findmax(identity, x, dims=dims)[2]
Base.argmax(itr) = findmax(identity, itr)[2]

and likewise for argmin. Whether generators are supported is somewhat orthogonal here.

[drewrobson]

Thanks very much for taking a look at this. I think you're right. Would the compiler be able to inline the identity function so there is no cost to this change in the most commonly used case?

[ararslan]

Yep. This is the same pattern we use in a lot of places where we define methods that accept function arguments.

[Nosferican]

Are there any other functions that could be prime candidates to get the function argument improvement? Besides, argmin, argmax... extrema?

[drewrobson]

Wait, I think there are some subtleties in reconciling @ararslan's nice suggestion with my original intent. Let me try to explain, and please forgive me if I confuse things further (e.g. I'm not entirely sure how to talk about the indices of an AbstractArray, but I gave it a shot!)

One of the difficulties in thinking about this is that there are really two flavors of argmax, the one defined in array.jl and the one defined in reducedim.jl.

In reducedim.jl, argmax defines an operation on A::AbstractArray. Here, I think the domain is implicitly CartesianIndices(axes(A)) and the function is implicitly getindex(A, i), where i ∈ CartesianIndices(axes(A)). The way I was thinking about it, argmax(A) is a convenient shorthand for something like argmax(i -> getindex(A, i), CartesianIndices(axes(A))) and I am proposing to generalize this to allow argmax(f, itr) for any function f and collection itr where this makes sense. In the special case that itr::AbstractArray then this would be implemented in reducedim.jl and support the dims kwarg and otherwise we have a fallback function in array.jl as long as keys(itr) is defined.

I think @ararslan's suggestion is distinct from this. If I understand correctly, for A::AbstractArray the domain would still implicitly be CartesianIndices(axes(A)) and the result of the new 2-arg flavor argmax(f, A) would not be an element of A (as proposed above), but rather an index of A. But please, correct me if I have misunderstood!

[ararslan]

We may be talking past each other, but I'll try to clarify.

For any iterable x, argmax(x) always returns an index into x, not an element of x. That can be seen more clearly from its maiden name, indmax. To ensure consistency with how the function is supposed to work, argmax(f, x) should return the index i at which f(x[i]) is maximized, not the element x[i]. To get the element, just index this back into x, e.g. x[argmax(f, x)].

[jekbradbury]

I think @ararslan’s is the standard Julian extension of a 1-argument function on collections to the case with an additional function argument. But @drewrobson is proposing the traditional 2-arg mathematical argmax, which is less consistent with the existing argmax but does coexist with it in math notation-land?

[drewrobson]

Yes exactly:

I completely agree that with the former name, it was clear that indmax should apply to things that can be indexed, and it should return the index that maximizes that thing. I'm also 100 % with you that if we want to stick with that original meaning even after the rename to argmax, then your suggestion is an undoubtably useful way to add functionality that is consistent with the original meaning of indmax.

With the rename to argmax, I am raising the (perhaps heretical!) suggestion that we should change what we mean by this function. To me, argmax now applies most naturally to a function, and it should return the arg that maximizes that function. As a special case, we would preserve the existing 1-arg flavor for things that can be indexed, with the understanding that getindex is the function to be maximized.

In summary, I think this is a backwards compatible and satisfying reinterpretation of what this function should do, but I get that it feels like a radical proposal!

[ararslan]

Ohhh okay, I see now, thanks @jekbradbury and @drewrobson for clarifying.

An unfortunate consequence of the renaming of indmax to argmax is indeed that the latter conflicts a bit with the identically named function from math when it comes to applying a function to the given collection. If we're to add a function argument to argmax I think we should ensure that the behavior is consistent with that of other Julia functions that accept function arguments, in particular that g(f, x) is equivalent to g(map(f, x)). I realize that's rather unsatisfying in this case given the more general name, but I really think making the cases behave differently will cause more confusion in the long run.

[drewrobson]

Yes, avoiding confusion is important. But then, perhaps argmax should only be defined in the 2-arg form proposed above, and then indmax should be defined as indmax(A) = argmax(i -> getindex(A, i), CartesianIndices(axes(A))) (or however you write that). Those names seem closest to what each function does, but I don't know if you'd want to split it up like that.

Edit: @Nosferican's post has reminded me that in addition to the 2-arg form proposed above, we would also keep the 1-arg argmax defined in array.jl for any itr for which keys/pairs is defined. Which will include generators, if #27612 happens!

[Nosferican]

• argmax is returning the index rather than the definition due to a historical accident (renaming of indmax to argmax) which should be fixed ASAP
• argmax works only on pairs-implemented collections due to the implementation which restricts the function to only a subset of all cases included in the mathematical definition (should be fixed)
• A 2-args method for passing a non identity function should be added
  My very naive implementation complies with the above,

iter = Set(1:3)
function my_argmax(f::Function, iter)
    next = iterate(iter)
    x = next[1]
    y = f(x)
    while next !== nothing
        (i, state) = next
        yᵢ = f(i)
        if y < yᵢ
            x = i[1]
        end
        next = iterate(iter, state)
    end
    return x
end

My ignorant self doesn't know what is the appropriate abstract type for an iterate.

[drewrobson]

@Nosferican, note that with #27612 this becomes a 1-liner. In other words, first define

Base.keys(g::Base.Generator) = g.iter

and then define

argmax(f::Function, itr) = argmax(f(x) for x in itr)

[ararslan]

argmax is returning the index rather than the definition due to a historical accident (renaming of indmax to argmax) which should be fixed ASAP

Can you elaborate on what you mean by that?