def parallel_fn(x1_or_kernel, x2=None, *args, **kwargs):
if utils.is_nt_tree_of(x1_or_kernel, np.ndarray):
return parallel_fn_x1(x1_or_kernel, x2, *args, **kwargs)
elif utils.is_nt_tree_of(x1_or_kernel, Kernel):
assert not x2
return parallel_fn_kernel(x1_or_kernel, *args, **kwargs)
raise NotImplementedError()
def serial_fn(x1_or_kernel: Union[NTTree[np.ndarray], NTTree[Kernel]],
x2: Optional[NTTree[Optional[np.ndarray]]] = None,
*args,
**kwargs) -> NTTree[Kernel]:
if utils.is_nt_tree_of(x1_or_kernel, np.ndarray):
return serial_fn_x1(x1_or_kernel, x2, *args, **kwargs)
elif utils.is_nt_tree_of(x1_or_kernel, Kernel):
if x2 is not None:
raise ValueError(f'`x2` must be `None`, got {x2}.')
return serial_fn_kernel(x1_or_kernel, *args, **kwargs)
else:
raise TypeError(x1_or_kernel, type(x1_or_kernel))