spine.utils.jit

JIT compilation decorators using Numba for performance optimization.

Functions

numbafy([cast_args, list_args, keep_torch, ...])

Function which wraps a numba function with some checks on the input to make the relevant conversions to numpy where necessary.
spine.utils.jit.numbafy(cast_args=[], list_args=[], keep_torch=False, ref_arg=None)[source]

Function which wraps a numba function with some checks on the input to make the relevant conversions to numpy where necessary.

Parameters:

  • cast_args (list(str), optional) – List of arguments to be cast to numpy

  • list_args (list(str), optional) – List of arguments which need to be cast to a numba typed list

  • keep_torch (bool, default False) – Make the output a torch object, if the reference argument is one

  • ref_arg (str, optional) – Reference argument used to assign a type and device to the torch output

Returns:

Wrapped function which ensures input type compatibility with numba

Return type:

callable