jax.profiler.trace_function

jax.profiler.trace_function(*args, **kwargs)