compute_flops and compute_vjp_flops for nnx.tabulate

[DBraun]

In Flax 0.11.1 it seems that nnx.tabulate doesn't support compute_flops and compute_vjp_flops like Linen did. I would like to track this feature or hear about alternatives.

[vfdev-5]

It was introduced in this PR and I think it could be possible to add these two args. I can work on that. We can do the following here:

flax/flax/nnx/summary.py

Lines 174 to 177 in 2bf5748

	rows.append(call_info)
	out = method(obj, *args, **kwargs)
	call_info.outputs = jax.tree.map(_to_dummy_array, out)
	return out

    rows.append(call_info)
    out = method(obj, *args, **kwargs)
    call_info.outputs = jax.tree.map(_to_dummy_array, out)
    if compute_flops:
      call_info.flops = _get_flops(method, obj, *args, **kwargs)
    
 

def _get_flops(fn, *args, **kwargs):
  e = jax.jit(fn).lower(*args, **kwargs).compile()
  cost = e.cost_analysis()
  if cost is None:
    return 0
  flops = int(cost['flops']) if 'flops' in cost else 0
  return flops

@cgarciae any feedback?

[cgarciae]

@vfdev-5 we do want this functionality. You're on the correct path, only modification I'd make is creating a wrapper function over fn so you can pass args and kwargs as a capture in case some of the inputs are static.