handle_jax_function

probly.representation.jax_functions.handle_jax_function(public_api: Callable[..., Any], relevant_args: Iterable[object], *args: object, **kwargs: object) → Any[source]

Dispatch a call to the __jax_function__ implementation of the arguments.

Parameters:

  • public_api – The wrapper function that was called, used as the dispatch key.

  • relevant_args – The arguments that may override the call.

  • *args – The positional arguments the wrapper was called with.

  • **kwargs – The keyword arguments the wrapper was called with.

Returns:

The first result that is not NotImplemented.

Raises:

TypeError – If no argument implements public_api.