Function _register_kernel_internal

torchvision/transforms/v2/functional/_utils.py:38–51

(functional, input_type, *, tv_tensor_wrapper=True)

def _register_kernel_internal(functional, input_type, *, tv_tensor_wrapper=True):
    registry = _KERNEL_REGISTRY.setdefault(functional, {})
    if input_type in registry:
        raise ValueError(f"Functional {functional} already has a kernel registered for type {input_type}.")

    def decorator(kernel):
        registry[input_type] = (
            _kernel_tv_tensor_wrapper(kernel)
            if issubclass(input_type, tv_tensors.TVTensor) and tv_tensor_wrapper
            else kernel
        )
        return kernel

    return decorator


def _name_to_functional(name):