Is there a generic way to count the number of calls to a function?

[marius311]

I want to do something like this to count how many times an objective function is called as part of some library which doesn't count things for me:

ncalls = 0

def foo(x):
    global ncalls
    ncalls = ncalls + 1
    return x**2

some_optimization_library(foo, x)

ncalls # <-- now look at this

Is there anything in this spirit which is possible in @jitted Jax (without modifying the optimization library)? Thanks.

[jakevdp]

You may be able to achieve what you have in mind using jax.experimental.host_callback. For example:

from jax.experimental import host_callback
from jax import jit

calls = 0

def register_call(arg, transforms):
  global calls
  calls += 1

@jit
def f(x):
  host_callback.id_tap(register_call, None)
  return x ** 2

@jit
def g(x):
  out = 0
  for i in range(5):
    out += f(x)
  return out

g(0.0)
print(calls)
# 5

Note however that this will likely result in performance degredation because of the required communication with the host.

[marius311]

Works perfectly, thanks! (and indeed only needed for some algorithmic benchmarks so performance hit doesnt matter)

[sharadmv]

Also consider the new jax.debug.callback (which will behave similarly to host_callback except it's not in jax.experimental.)

[marius311]

@sharadmv Many thanks for that. Actually once I used the host_callback solution from above in some less-trivial code, I was getting the error below. However, the less-trivial code works with jax.debug.callback.

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
/tmp/ipykernel_94908/3444581130.py in <cell line: 6>()
      4     prob.get_J(result)
      5     prob.get_H(result, nsims=10, progress=True)
----> 6 ncalls_muse = count_calls(solve)

/tmp/ipykernel_94908/3187089758.py in count_calls(f)
      8     global ncalls
      9     ncalls = 0
---> 10     f()
     11     return ncalls

/tmp/ipykernel_94908/3444581130.py in solve()
      3     prob.solve(θ_start=θ_true, result=result, progress=True, get_covariance=False)
      4     prob.get_J(result)
----> 5     prob.get_H(result, nsims=10, progress=True)
      6 ncalls_muse = count_calls(solve)

~/work/muse-inference/muse_inference/muse_inference.py in get_H(self, result, θ, step, method, θ_tol, z_tol, rng, nsims, pmap, progress, skip_errors)
    496             rngs = self._split_rng(rng, nsims_remaining)
    497             _get_H_i = partial(self._get_H_i, θ=θ, method=method, θ_tol=θ_tol, z_tol=z_tol, step=step, skip_errors=skip_errors)
--> 498             result.Hs.extend(H for H in pbar(map(_get_H_i, rngs)) if H is not None)
    499             result.time += datetime.now() - t0
    500 

~/work/muse-inference/muse_inference/muse_inference.py in <genexpr>(.0)
    496             rngs = self._split_rng(rng, nsims_remaining)
    497             _get_H_i = partial(self._get_H_i, θ=θ, method=method, θ_tol=θ_tol, z_tol=z_tol, step=step, skip_errors=skip_errors)
--> 498             result.Hs.extend(H for H in pbar(map(_get_H_i, rngs)) if H is not None)
    499             result.time += datetime.now() - t0
    500 

~/work/muse-inference/.venv/lib/python3.9/site-packages/tqdm/std.py in __iter__(self)
   1193 
   1194         try:
-> 1195             for obj in iterable:
   1196                 yield obj
   1197                 # Update and possibly print the progressbar.

    [... skipping hidden 11 frame]

~/work/muse-inference/.venv/lib/python3.9/site-packages/jax/experimental/host_callback.py in <lambda>(j)
   1814 xla.register_translation(id_p, lambda ctx, avals_in, avals_out, *args: args)
   1815 
-> 1816 dispatch.outfeed_rewriter = lambda j: _rewrite_jaxpr(j, False, False)
   1817 
   1818 

~/work/muse-inference/.venv/lib/python3.9/site-packages/jax/experimental/host_callback.py in _rewrite_jaxpr(jaxpr, has_input_token, has_output_token)
   1500       output_token_var = mk_new_var(last_token_var.aval)
   1501       output_itoken_var = mk_new_var(last_itoken_var.aval)
-> 1502       _rewrite_eqn(eqn, eqns, last_token_var, output_token_var,
   1503                    last_itoken_var, output_itoken_var, mk_new_var)
   1504       last_token_var = output_token_var

~/work/muse-inference/.venv/lib/python3.9/site-packages/jax/experimental/host_callback.py in _rewrite_eqn(eqn, eqns, input_token_var, output_token_var, input_itoken_var, output_itoken_var, mk_new_var)
   1697             )))
   1698   else:
-> 1699     raise NotImplementedError(f"outfeed rewrite {eqn.primitive}")
   1700 
   1701 

NotImplementedError: outfeed rewrite custom_linear_solve