SupportsJaxFunction

class probly.representation.jax_functions.SupportsJaxFunction(*args, **kwargs)

Bases: Protocol

Structural interface for custom JAX function overrides.

An override may return a representation, an array, or another result type. This interface deliberately does not promise that operations preserve type.