Random expensive recompilations with BFGS

[mattwescott]

The same optimization with the same input shapes sometimes takes a few seconds, and sometimes it is 1000x faster. Maybe JIT is recompiling based on dynamic structure in BFGS? I'm seeing random slow evaluations even with maxiter==1, line_search_maxiter==1.

Is this expected when the number of iterations is fixed? If it needs digging, I'm happy to work on a minimal reproduction.

@jax.jit
def optimize(f, x0):
    return opt.minimize(
          loss_fn,
          x0=x0,
          method='BFGS',
          options=dict(maxiter=1, line_search_maxiter=1, gtol=1e-2),
      )

[jakevdp]

A minimal repro would probably help. One basic question: are you timing using .block_until_ready()? If not, your fast runs might be just measuring the dispatch time (see Asynchronous Dispatch for more info on running microbenchmarks).

[mattwescott]

Thanks Jake, this was my mistake. I had the wrong static_argnums in a jit compile

[hawkinsp]

By the way, are you aware that static_argnames is a thing on jit these days? Names can be easier to get right!