jaxrts.instrument_function.FromCallable.__init__

FromCallable.__init__(function, disable_partial_unwrapping: bool = False)

Create an :py:class`~.InstrumentFunction` from a callable.

If function is a functools.partial, its bound positional and keyword arguments are unpacked and re-flattened as genuine pytree leaves. This requires that tracing though all positional and keyword arguments is possible. If this is not the case, define a new function where this is the case, or set the flag “disable_partial_unwrapping==True”.