Most efficient way to find argmin and min of f(array)

[aeftimia]

Say I have an array of inputs x for which I would like to find the lowest result of applying f and the index of that minimal value.

Something like this:

lowest = float('inf')
for i, e in enumerate(x):
    new = f(e)
    if new < lowest:
        lowest = new
        lowest_i = i

I think the direct equivalent in jax would be a while_loop but is there any way to parallelize this process? I'd imagine you could also use vmap or pmap followed by argmin, but that of course allocates memory for all executions of f.

How would you recommend implementing this process in jax?

[soraros]

Wouldn't jnp.argmin(f(x)) work?

[aeftimia]

This only works if the minimum is to be taken with respect to all axes of x. In my case, I just want the first, and would also like to retrieve the value of that minimum. But even in the case of a one dimensional x, you're storing all the results of f(x) before identifying the minimum... This is O(N) in space. It's clearly possible to produce a O(1) space procedure that's sequential. The flip side is that once you have all the results of f(x) stored, you can do efficient parallel procedures that search for the minimum (and its location).

I'm getting the sense there would need to be some kind of XLA primitive that combines map and reduce in such a way that f is mapped to the value in the array before passing the result off to reduce.

[soraros]

argmin does have an axis keyword arguments.

And btw, I think your particular needs are worth mentioning in the question itself.

[aeftimia]

Searching around for a "mapreduce" primitive in jax lead me to this thread. Looks like the answer is in fact to use vmap with the reduction applied, though apparently scan can work too under some circumstances.... But in any case jit is crucial to optimize the XLA computations.

[jakevdp]

I would suggest using the most straightforward approach:

@jit
def argmin_and_min(x):
  y = f(x)
  i = jnp.argmin(y)
  return i, y[i]

lowest_i, lowest_y = argmin_and_min(x)

You ask "is there any way to parallelize this process?"; if you're running this on an accelerator like GPU or TPU, the computation will be parallelized automatically within the hardware. While it may be tempting to try to find a solution based on while_loop, scan, or other control flow operations, this will generally lead to less efficient computation when running on accelerators when compared to these native vectorized operations.

Above you mentioned being concerned about the need to allocate the full array y in this approach. While this concern is valid when running in op-by-op mode (i.e. outside jit), I wouldn't worry about it a priori when constructing what looks like a temporary variable within a JIT-compiled function definition. When compiled by XLA, the compiler will generally fuse operations in order to compute equivalent results while avoiding unnecessary operations or memory allocation.