mlreco.utils.numba module¶
-
mlreco.utils.numba.numba_wrapper(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 ([str]) – List of arguments to be cast to numpy
list_args ([str]) – List of arguments which need to be cast to a numba typed list
keep_torch (bool) – Make the output a torch object, if the reference argument is one
ref_arg (str) – Reference argument used to assign a type and device to the torch output
- Returns
Function
-
mlreco.utils.numba.unique_nb(x: numba.int32.slice(None, None, None))¶
-
mlreco.utils.numba.submatrix_nb(x: numba.float32.(slice(None, None, None), slice(None, None, None)), index1: numba.int32.slice(None, None, None), index2: numba.int32.slice(None, None, None)) -> numba.float32.(slice(None, None, None), slice(None, None, None))¶ Numba implementation of matrix subsampling
-
mlreco.utils.numba.cdist_nb(x1: numba.float32.(slice(None, None, None), slice(None, None, None)), x2: numba.float32.(slice(None, None, None), slice(None, None, None))) -> numba.float32.(slice(None, None, None), slice(None, None, None))¶ Numba implementation of Eucleadian cdist in 3D.
-
mlreco.utils.numba.mean_nb(x: numba.float32.(slice(None, None, None), slice(None, None, None)), axis: numba.int32) -> numba.float32.slice(None, None, None)¶ Numba implementation of np.mean(x, axis)
-
mlreco.utils.numba.argmin_nb(x: numba.float32.(slice(None, None, None), slice(None, None, None)), axis: numba.int32) -> numba.int32.slice(None, None, None)¶ Numba implementation of np.argmin(x, axis)
-
mlreco.utils.numba.argmax_nb(x: numba.float32.(slice(None, None, None), slice(None, None, None)), axis: numba.int32) -> numba.int32.slice(None, None, None)¶ Numba implementation of np.argmax(x, axis)
-
mlreco.utils.numba.min_nb(x: numba.float32.(slice(None, None, None), slice(None, None, None)), axis: numba.int32) -> numba.float32.slice(None, None, None)¶ Numba implementation of np.max(x, axis)
-
mlreco.utils.numba.max_nb(x: numba.float32.(slice(None, None, None), slice(None, None, None)), axis: numba.int32) -> numba.float32.slice(None, None, None)¶ Numba implementation of np.max(x, axis)
-
mlreco.utils.numba.all_nb(x: numba.float32.(slice(None, None, None), slice(None, None, None)), axis: numba.int32) -> numba.int32.slice(None, None, None)¶ Numba implementation of np.all(x, axis)
-
mlreco.utils.numba.softmax_nb(x: numba.float32.(slice(None, None, None), slice(None, None, None)), axis: numba.int32) -> numba.float32.(slice(None, None, None), slice(None, None, None))¶
-
mlreco.utils.numba.log_loss_nb(x1: numba.boolean.slice(None, None, None), x2: numba.float32.slice(None, None, None)) → numba.float32¶