jit

seli.jit(function: Callable[[P], T]) → Callable[[P], T]

Just-in-time compiling functions.

This is a drop-in replacement for jax.jit, that traces shared references between the different arguments. Using jax.jit with references shared between arguments will untie the references in the body of the function and the output.

The function will not recompile if only the values inside of the jax.Arrays change.

Parameters:

function (Callable) – Function to apply the jax.jit to.

Returns:

compiled – The compiled function. This function takes the same arguments as the original function.

Return type:

Callable