def jit(signature_or_function=None, **kwargs): if signature_or_function is None: return lambda x: x else: return signature_or_function