try_jax_function

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

Try function overrides without preventing Python’s reflected operator fallback.

Parameters:

• public_api – Wrapper function used as the dispatch key.

• relevant_args – Arguments whose types may override the function.

• *args – Positional arguments for the function.

• **kwargs – Keyword arguments for the function.

Returns:

The first implemented result, or NotImplemented if every override declines. Exceptions raised by implementations propagate unchanged.