has_jax_function

probly.representation.jax_functions.has_jax_function(relevant_args: Iterable[object]) → bool

Check whether any of the arguments overrides the jax function protocol.

Parameters:

relevant_args – The arguments to inspect. None entries are ignored.

Returns:

True if the type of at least one argument defines __jax_function__.