import sys
_module = sys.modules[__name__]
del sys
extensions = _module
stability = _module
conf = _module
__about__ = _module
flash = _module
audio = _module
classification = _module
cli = _module
data = _module
input = _module
input_transform = _module
speech_recognition = _module
backbone = _module
collate = _module
data = _module
input = _module
model = _module
output_transform = _module
core = _module
adapter = _module
classification = _module
base_viz = _module
batch = _module
callback = _module
data_module = _module
io = _module
classification_input = _module
input = _module
output = _module
transform_predictions = _module
properties = _module
splits = _module
transforms = _module
utilities = _module
classification = _module
collate = _module
data_frame = _module
loading = _module
paths = _module
samples = _module
sort = _module
utils = _module
finetuning = _module
heads = _module
hooks = _module
integrations = _module
fiftyone = _module
icevision = _module
adapter = _module
backbones = _module
transforms = _module
wrappers = _module
labelstudio = _module
input = _module
visualizer = _module
pytorch_forecasting = _module
adapter = _module
transforms = _module
pytorch_tabular = _module
transformers = _module
collate = _module
model = _module
optimizers = _module
lamb = _module
lars = _module
lr_scheduler = _module
optimizers = _module
schedulers = _module
registry = _module
regression = _module
serve = _module
_compat = _module
cached_property = _module
component = _module
composition = _module
core = _module
dag = _module
optimization = _module
order = _module
rewrite = _module
task = _module
utils_test = _module
visualize = _module
decorators = _module
execution = _module
flash_components = _module
interfaces = _module
http = _module
models = _module
templates = _module
server = _module
types = _module
base = _module
bbox = _module
image = _module
label = _module
number = _module
repeated = _module
table = _module
text = _module
trainer = _module
apply_func = _module
compatibility = _module
embedder = _module
flash_cli = _module
imports = _module
isinstance = _module
lightning_cli = _module
providers = _module
stages = _module
types = _module
url_error = _module
graph = _module
data = _module
input = _module
model = _module
collate = _module
embedding = _module
model = _module
adapters = _module
clip = _module
resnet = _module
timm = _module
torchvision = _module
transformers = _module
data = _module
input_transform = _module
baal = _module
data = _module
dropout = _module
loop = _module
learn2learn = _module
model = _module
data = _module
detection = _module
data = _module
vissl_heads = _module
losses = _module
vissl_losses = _module
model = _module
strategies = _module
default = _module
vissl_strategies = _module
vissl_transforms = _module
vissl = _module
adapter = _module
hooks = _module
multicrop = _module
utilities = _module
face_detection = _module
fastface_backbones = _module
data = _module
input = _module
input_transform = _module
model = _module
instance_segmentation = _module
data = _module
keypoint_detection = _module
segmentation = _module
data = _module
heads = _module
model = _module
output = _module
viz = _module
style_transfer = _module
data = _module
model = _module
pointcloud = _module
data = _module
datasets = _module
input = _module
model = _module
open3d_ml = _module
app = _module
backbones = _module
data = _module
input = _module
model = _module
app = _module
backbones = _module
sequences_dataset = _module
setup_tools = _module
tabular = _module
model = _module
forecasting = _module
data = _module
model = _module
template = _module
backbones = _module
data = _module
model = _module
adapters = _module
clip = _module
huggingface = _module
model = _module
ort_callback = _module
question_answering = _module
model = _module
seq2seq = _module
model = _module
summarization = _module
model = _module
translation = _module
video = _module
data = _module
input = _module
input_transform = _module
model = _module
utils = _module
audio_classification = _module
face_detection = _module
graph_classification = _module
graph_embedder = _module
image_classification = _module
image_classification_multi_label = _module
image_embedder = _module
image_classification_active_learning = _module
image_classification = _module
image_classification_fiftyone_datasets = _module
image_embedding = _module
object_detection = _module
text_classification = _module
video_classification = _module
image_classification_imagenette_mini = _module
tabular_forecasting_interpretable = _module
pointcloud_detection = _module
pointcloud_segmentation = _module
semantic_segmentation = _module
client = _module
inference_server = _module
inference = _module
speech_recognition = _module
style_transfer = _module
tabular_classification = _module
tabular_forecasting = _module
tabular_regression = _module
template = _module
text_classification = _module
text_classification_multi_label = _module
text_embedder = _module
translation = _module
video_classification = _module
pointcloud_detection = _module
pointcloud_segmentation = _module
setup = _module
tests = _module
test_data = _module
test_model = _module
test_data_model_integration = _module
test_model = _module
conftest = _module
test_input = _module
test_output = _module
test_output_transform = _module
test_base_viz = _module
test_batch = _module
test_callback = _module
test_callbacks = _module
test_data_module = _module
test_input_transform = _module
test_properties = _module
test_splits = _module
test_transforms = _module
test_classification = _module
test_loading = _module
test_paths = _module
test_labelstudio = _module
test_strategies = _module
test_lr_scheduler = _module
test_optimizers = _module
models = _module
test_compat = _module
test_cached_property = _module
test_components = _module
test_composition = _module
test_dag = _module
test_optimization = _module
test_order = _module
test_rewrite = _module
test_task = _module
test_utils = _module
test_gridbase_validations = _module
test_integration = _module
test_types = _module
test_bbox = _module
test_image = _module
test_label = _module
test_number = _module
test_repeated = _module
test_table = _module
test_text = _module
test_classification = _module
test_data = _module
test_finetuning = _module
test_model = _module
test_registry = _module
test_trainer = _module
test_embedder = _module
test_lightning_cli = _module
test_stability = _module
deprecated_api = _module
test_remove_0_9_0 = _module
examples = _module
test_integrations = _module
test_scripts = _module
test_model = _module
test_model = _module
helpers = _module
boring_model = _module
task_tester = _module
test_active_learning = _module
test_data = _module
test_model = _module
test_training_strategies = _module
test_data_model_integration = _module
test_model = _module
test_output = _module
test_model = _module
test_data = _module
test_model = _module
test_model = _module
test_backbones = _module
test_data = _module
test_heads = _module
test_model = _module
test_output = _module
test_model = _module
test_data = _module
test_data = _module
test_datasets = _module
test_model = _module
test_model = _module
test_model = _module
test_model = _module
test_model = _module
test_model = _module
test_model = _module
test_model = _module
test_model = _module
test_model = _module
tpu = _module
test_sample_tpu = _module
test_tpu_multi_core = _module
test_tpu_single_core = _module
test_data = _module
test_model = _module

from _paritybench_helpers import _mock_config, patch_functional
from unittest.mock import mock_open, MagicMock
from torch.autograd import Function
from torch.nn import Module
import abc, collections, copy, enum, functools, inspect, itertools, logging, math, matplotlib, numbers, numpy, pandas, queue, random, re, scipy, sklearn, string, tensorflow, time, torch, torchaudio, torchtext, torchvision, types, typing, uuid, warnings
import numpy as np
from torch import Tensor
patch_functional()
open = mock_open()
yaml = logging = sys = argparse = MagicMock()
ArgumentParser = argparse.ArgumentParser
_global_config = args = argv = cfg = config = params = _mock_config()
argparse.ArgumentParser.return_value.parse_args.return_value = _global_config
yaml.load.return_value = _global_config
sys.argv = _global_config
__version__ = '1.0.0'
xrange = range
wraps = functools.wraps


import warnings


from typing import Any


from typing import Callable


from typing import Collection


from typing import Dict


from typing import List


from typing import Optional


from typing import Sequence


from typing import Type


from typing import Union


import numpy as np


import pandas as pd


from torch import Tensor


from typing import Tuple


import torch


from torch.utils.data import Dataset


from typing import Mapping


import torch.nn as nn


from abc import abstractmethod


import torch.jit


from torch import nn


from torch.utils.data import DataLoader


from torch.utils.data import Sampler


import torch.nn.functional as F


from typing import TYPE_CHECKING


from typing import Iterable


from torch.utils.data.dataset import IterableDataset


from torch.utils.data.sampler import Sampler


import functools


from enum import Enum


from typing import cast


from functools import reduce


from typing import ClassVar


from torch.utils.data._utils.collate import default_collate as torch_default_collate


import copy


from functools import partial


from typing import Set


from torch.nn import Module


from torch.optim import Optimizer


from inspect import getmembers


from copy import deepcopy


from copy import copy


from torch.utils.data._utils.collate import default_collate


import inspect


import re


from abc import ABCMeta


from torch.optim.lr_scheduler import _LRScheduler


from torch.optim.optimizer import Optimizer


import math


from torch.optim.optimizer import required


from torch.optim import Adam


from inspect import isclass


from torch import optim


from torch.optim import lr_scheduler


from torch.optim.lr_scheduler import CosineAnnealingLR


from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts


from torch.optim.lr_scheduler import CyclicLR


from torch.optim.lr_scheduler import MultiStepLR


from torch.optim.lr_scheduler import ReduceLROnPlateau


from torch.optim.lr_scheduler import StepLR


from functools import wraps


from typing import Iterator


from typing import NamedTuple


from typing import TypeVar


import types


from types import MethodType


from torch.nn import functional as F


from torch.nn import Linear


from torch.utils.data.dataloader import default_collate


from typing import IO


from collections import defaultdict


from torch.utils.data import IterableDataset


from torch.hub import load_state_dict_from_url


from torch.utils.data import random_split


from torch.utils.data._utils.worker import get_worker_info


from types import FunctionType


from torch import tensor


import random


from typing import NoReturn


from torch.utils.data.dataset import Dataset


from abc import ABC


import logging


import collections


from torch.utils.data import DistributedSampler


from torchvision.datasets import CIFAR10


from itertools import chain


import torchvision.transforms as T


import matplotlib.pyplot as plt


import sklearn.datasets


from sklearn import datasets


from collections import namedtuple


import time


from numbers import Number


from torch.nn import Flatten


from torch.nn import LogSoftmax


from torch.utils.data import Subset


from torch.utils.data import SequentialSampler


from pandas import DataFrame


def _has_len(data: Union[Sequence, Iterable]) ->bool:
    """Duck typing check to see if the argument supports getting the length.

    Args:
        data: The object to check for length support.
    """
    try:
        len(data)
        return True
    except (TypeError, NotImplementedError):
        return False


def _validate_input(input: 'InputBase') ->None:
    """Helper function to validate that the type of an ``InputBase.data`` is appropriate for the type of
    ``InputBase`` being used.

    Args:
        input: The ``InputBase`` instance to validate.

    Raises:
        RuntimeError: If the ``input`` is of type ``Input`` and it's ``data`` attribute does not support ``len``.
        RuntimeError: If the ``input`` is of type ``IterableInput`` and it's ``data`` attribute does support ``len``.
    """
    if input.data is not None:
        if isinstance(input, Input) and not _has_len(input.data):
            raise RuntimeError('`Input.data` is not a sequence with a defined length. Use `IterableInput` instead.')
        elif isinstance(input, IterableInput) and _has_len(input.data):
            raise RuntimeError('`IterableInput.data` is a sequence with a defined length. Use `Input` instead.')


def _wrap_init(class_dict: Dict[str, Any]) ->None:
    """Helper function to wrap the ``__init__`` (if present) from a class construction dict to apply the
    ``_validate_input`` function after instantiation. Modifies the dict inplace.

    Args:
        class_dict: The class construction dict, optionally containing an init to wrap.
    """
    if '__init__' in class_dict:
        fn = class_dict['__init__']

        @functools.wraps(fn)
        def wrapper(self, *args, **kwargs):
            fn(self, *args, **kwargs)
            _validate_input(self)
        class_dict['__init__'] = wrapper


def _deepcopy_dict(nested_dict: Any) ->Any:
    """Utility to deepcopy a nested dict."""
    if not isinstance(nested_dict, Dict):
        return nested_dict
    return {key: value for key, value in nested_dict.items()}


PROVIDERS = []


_REGISTERED_FUNCTION = Dict[str, Any]


def print_provider_info(name, providers, func):
    if not isinstance(providers, List):
        providers = [providers]
    providers = list(providers)
    if len(providers) > 1:
        providers[-2] = f'{str(providers[-2])} and {str(providers[-1])}'
        providers = providers[:-1]
    message = f"Using '{name}' provided by {', '.join(str(provider) for provider in providers)}."

    def build_wrapper(func):

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            rank_zero_info(message)
            return func(*args, **kwargs)
        return wrapper
    wrapper = build_wrapper(func)
    if inspect.isclass(func):
        callables = [f for f in dir(func) if callable(getattr(func, f)) and not f.startswith('_')]
        for c in callables:
            setattr(wrapper, c, build_wrapper(getattr(func, c)))
    return wrapper


def _is_list_like(x: Any) ->bool:
    try:
        _ = x[0]
        _ = len(x)
        return True
    except (TypeError, IndexError, KeyError):
        return False


def _is_list_like_excluding_str(x):
    return _is_list_like(x) and str(x) != x


def default_uncollate(batch: Any) ->List[Any]:
    """This function is used to uncollate a batch into samples. The following conditions are used:

    - if the ``batch`` is a ``dict``, the result will be a list of dicts
    - if the ``batch`` is list-like, the result is guaranteed to be a list

    Args:
        batch: The batch of outputs to be uncollated.

    Returns:
        The uncollated list of predictions.

    Raises:
        ValueError: If the input is a ``dict`` whose values are not all list-like.
        ValueError: If the input is a ``dict`` whose values are not all the same length.
        ValueError: If the input is not a ``dict`` or list-like.
    """
    if isinstance(batch, dict):
        if any(not _is_list_like_excluding_str(sub_batch) for sub_batch in batch.values()):
            raise ValueError('When uncollating a dict, all sub-batches (values) are expected to be list-like.')
        if len({len(sub_batch) for sub_batch in batch.values()}) > 1:
            raise ValueError('When uncollating a dict, all sub-batches (values) are expected to have the same length.')
        elements = [default_uncollate(element) for element in zip(*batch.values())]
        return [dict(zip(batch.keys(), element)) for element in elements]
    if isinstance(batch, (list, tuple, Tensor)):
        return list(batch)
    raise ValueError(f'The batch of outputs to be uncollated is expected to be a `dict` or list-like (e.g. `Tensor`, `list`, `tuple`, etc.), but got input of type: {type(batch)}')


class OutputTransform:
    """The :class:`~flash.core.data.io.output_transform.OutputTransform` encapsulates all the data processing logic
    that should run after the model."""

    @staticmethod
    def per_batch_transform(batch: Any) ->Any:
        """Transforms to apply on a whole batch before uncollation to individual samples.

        Can involve both CPU and Device transforms as this is not applied in separate workers.
        """
        return batch

    @staticmethod
    def per_sample_transform(sample: Any) ->Any:
        """Transforms to apply to a single sample after splitting up the batch.

        Can involve both CPU and Device transforms as this is not applied in separate workers.
        """
        return sample

    @staticmethod
    def uncollate(batch: Any) ->Any:
        """Uncollates a batch into single samples.

        Tries to preserve the type wherever possible.
        """
        return default_uncollate(batch)

    def __call__(self, batch: Sequence[Any]):
        if batch is None:
            return batch
        uncollated = self.uncollate(self.per_batch_transform(batch))
        return [self.per_sample_transform(sample) for sample in uncollated]


def requires(*module_paths: Union[str, Tuple[bool, str]]):

    def decorator(func):
        available = True
        extras = []
        modules = []
        for module_path in module_paths:
            if isinstance(module_path, str):
                if module_path in _EXTRAS_AVAILABLE:
                    extras.append(module_path)
                    if not _EXTRAS_AVAILABLE[module_path]:
                        available = False
                else:
                    modules.append(module_path)
                    if not module_available(module_path):
                        available = False
            else:
                available, module_path = module_path
                modules.append(module_path)
        if not available:
            modules = [f"'{module}'" for module in modules]
            if extras:
                modules.append(f"'lightning-flash[{','.join(extras)}]'")

            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                raise ModuleNotFoundError(f"Required dependencies not available. Please run: pip install {' '.join(modules)}")
            return wrapper
        return func
    return decorator


class CheckDependenciesMeta(ABCMeta):

    def __new__(mcs, *args, **kwargs):
        result = ABCMeta.__new__(mcs, *args, **kwargs)
        if result.required_extras is not None:
            if isinstance(result.required_extras, str):
                result.required_extras = [result.required_extras]
            result.__init__ = requires(*result.required_extras)(result.__init__)
            patterns = ['load_from_checkpoint', 'available_*']
            regex = '(' + ')|('.join(patterns) + ')'
            for attribute_name, attribute_value in filter(lambda x: re.match(regex, x[0]), inspect.getmembers(result)):
                if attribute_name in ['available_layers']:
                    continue
                setattr(result, attribute_name, classmethod(requires(*result.required_extras)(attribute_value.__func__)))
        return result


class Connection(NamedTuple):
    """A connection maps one output to one input.

    This is a self-contained data structure, which when given in the context of
    the other components in a composition, will map input/output keys/indices
    between components.

    Warnings
    --------
    * This data structure should not be instantiated directly! The
      class_methods attached to the class are the indended mechanisms to create
      a new instance.
    """
    source_component: str
    target_component: str
    source_key: str
    target_key: str

    def __repr__(self):
        return f'Connection({str(self)})'

    def _repr_pretty_(self, p, cycle):
        if cycle:
            return
        res = f'Connection({self.source_component}.outputs.{self.source_key} >> {self.target_component}.inputs.{self.target_key})'
        p.text(res)

    def __str__(self):
        return f'{self.source_component}.outputs.{self.source_key} >> {self.target_component}.inputs.{self.target_key}'


def flatten(seq, container=list):
    """
    >>> list(flatten([1]))
    [1]
    >>> list(flatten([[1, 2], [1, 2]]))
    [1, 2, 1, 2]
    >>> list(flatten([[[1], [2]], [[1], [2]]]))
    [1, 2, 1, 2]
    >>> list(flatten(((1, 2), (1, 2)))) # Don't flatten tuples
    [(1, 2), (1, 2)]
    >>> list(flatten((1, 2, [3, 4]))) # support heterogeneous
    [1, 2, 3, 4]
    """
    if isinstance(seq, str):
        yield seq
    else:
        for item in seq:
            if isinstance(item, container):
                yield from flatten(item, container=container)
            else:
                yield item


def connections_from_components_map(components: Dict[str, 'ModelComponent']) ->List[Dict[str, str]]:
    dsk_connections = []
    for con in flatten([comp._flashserve_meta_.connections for comp in components.values()]):
        dsk_connections.append(con._asdict())
    return dsk_connections


def funcname(func):
    """Get the name of a function."""
    if isinstance(func, functools.partial):
        return funcname(func.func)
    if isinstance(func, methodcaller):
        return str(func)[:50]
    module_name = getattr(func, '__module__', None) or ''
    type_name = getattr(type(func), '__name__', None) or ''
    if 'cytoolz' in module_name and type_name == 'curry':
        return func.func_name[:50]
    if 'numpy' in module_name and type_name == 'vectorize':
        return ('vectorize_' + funcname(func.pyfunc))[:50]
    try:
        name = func.__name__
        if name == '<lambda>':
            return 'lambda'
        return name[:50]
    except AttributeError:
        return str(func)[:50]


def istask(x):
    """Is x a runnable task?
    A task is a tuple with a callable first argument
    Examples
    --------
    >>> istask((inc, 1))
    True
    >>> istask(1)
    False
    """
    return type(x) is tuple and x and callable(x[0])


def unwrap_partial(func):
    while hasattr(func, 'func'):
        func = func.func
    return func


def functions_of(task):
    """Set of functions contained within nested task.

    Examples
    --------
    >>> task = (add, (mul, 1, 2), (inc, 3))  # doctest: +SKIP
    >>> functions_of(task)  # doctest: +SKIP
    set([add, mul, inc])
    """
    funcs = set()
    work = [task]
    sequence_types = {list, tuple}
    while work:
        new_work = []
        for task in work:
            if type(task) in sequence_types:
                if istask(task):
                    funcs.add(unwrap_partial(task[0]))
                    new_work.extend(task[1:])
                else:
                    new_work.extend(task)
        work = new_work
    return funcs


no_default = '__no_default__'


def get_dependencies(dsk, key=None, task=no_default, as_list=False):
    """Get the immediate tasks on which this task depends.

    Examples
    --------
    >>> dsk = {'x': 1,
    ...        'y': (inc, 'x'),
    ...        'z': (add, 'x', 'y'),
    ...        'w': (inc, 'z'),
    ...        'a': (add, (inc, 'x'), 1)}
    >>> get_dependencies(dsk, 'x')
    set()
    >>> get_dependencies(dsk, 'y')
    {'x'}
    >>> get_dependencies(dsk, 'z')  # doctest: +SKIP
    {'x', 'y'}
    >>> get_dependencies(dsk, 'w')  # Only direct dependencies
    {'z'}
    >>> get_dependencies(dsk, 'a')  # Ignore non-keys
    {'x'}
    >>> get_dependencies(dsk, task=(inc, 'x'))  # provide tasks directly
    {'x'}
    """
    if key is not None:
        arg = dsk[key]
    elif task is not no_default:
        arg = task
    else:
        raise ValueError('Provide either key or task')
    result = []
    work = [arg]
    while work:
        new_work = []
        for w in work:
            typ = type(w)
            if typ is tuple and w and callable(w[0]):
                new_work.extend(w[1:])
            elif typ is list:
                new_work.extend(w)
            elif typ is dict:
                new_work.extend(w.values())
            else:
                try:
                    if w in dsk:
                        result.append(w)
                except TypeError:
                    pass
        work = new_work
    return result if as_list else set(result)


def reverse_dict(d):
    """
    >>> a, b, c = 'abc'
    >>> d = {a: [b, c], b: [c]}
    >>> reverse_dict(d)  # doctest: +SKIP
    {'a': set([]), 'b': set(['a']}, 'c': set(['a', 'b'])}
    """
    result = defaultdict(set)
    _add = set.add
    for k, vals in d.items():
        result[k]
        for val in vals:
            _add(result[val], k)
    result.default_factory = None
    return result


def get_deps(dsk):
    """Get dependencies and dependents from task graph.

    Examples
    --------
    >>> dsk = {'a': 1, 'b': (inc, 'a'), 'c': (inc, 'b')}
    >>> dependencies, dependents = get_deps(dsk)
    >>> dependencies
    {'a': set(), 'b': {'a'}, 'c': {'b'}}
    >>> dict(dependents)
    {'a': {'b'}, 'b': {'c'}, 'c': set()}
    """
    dependencies = {k: get_dependencies(dsk, task=v) for k, v in dsk.items()}
    dependents = reverse_dict(dependencies)
    return dependencies, dependents


def component_dag_content(components: Dict[str, 'ModelComponent']) ->'ComponentJSON':
    dsk_connections = connections_from_components_map(components)
    comp_dependencies, comp_dependents, comp_funcnames = {}, {}, {}
    for comp_name, comp in components.items():
        functions_comp = valmap(functions_of, comp._flashserve_meta_.dsk)
        function_names_comp = {k: sorted(set(map(funcname, v))) for k, v in functions_comp.items()}
        comp_funcnames[comp_name] = function_names_comp
        _dependencies, _dependents = get_deps(comp._flashserve_meta_.dsk)
        _dependents = dict(_dependents)
        comp_dependencies[comp_name] = _dependencies
        comp_dependents[comp_name] = _dependents
    return ComponentJSON(component_dependencies=comp_dependencies, component_dependents=comp_dependents, component_funcnames=comp_funcnames, connections=dsk_connections)


def endpoint_protocol_content(ep_proto: 'EndpointProtocol') ->'EndpointProtoJSON':
    ep_proto_payload_dsk_key_map = valmap(lambda x: f'{x}.serial', ep_proto.dsk_input_key_map)
    ep_proto_result_key_dsk_map = valmap(lambda x: f'{x}.serial', ep_proto.dsk_output_key_map)
    return EndpointProtoJSON(name=ep_proto.name, route=ep_proto.route, payload_key_dsk_task=ep_proto_payload_dsk_key_map, result_key_dsk_task=ep_proto_result_key_dsk_map)


def merged_dag_content(ep_proto: 'EndpointProtocol', components: Dict[str, 'ModelComponent']) ->'MergedJSON':
    init = _process_initial(ep_proto, components)
    dsk_connections = connections_from_components_map(components)
    epjson = endpoint_protocol_content(ep_proto)
    merged = {**init.merged_dsk, **init.payload_tasks_dsk}
    dependencies, _ = get_deps(merged)
    merged_proto = defaultdict(list)
    for task_name, task in merged.items():
        for parent in dependencies[task_name]:
            merged_proto[task_name].append(parent)
    for request_name, task_key in init.payload_dsk_map.items():
        cluster, *_ = task_key.split('.')
        merged_proto[task_key[:-len('.serial')]].append(task_key)
        merged_proto[task_key].append(request_name)
    merged_proto = dict(merged_proto)
    dependencies, dependents = get_deps(merged_proto)
    dependents = dict(dependents)
    functions_merged = valmap(functions_of, merged)
    function_names_merged = {k: tuple(map(funcname, v)) for k, v in functions_merged.items()}
    return MergedJSON(dependencies=dependencies, dependents=dependents, funcnames=function_names_merged, connections=dsk_connections, endpoint=epjson)


def _build_dag_json(components: Dict[str, 'ModelComponent'], ep_proto: Optional['EndpointProtocol'], *, show_connected_components: bool=True):
    if show_connected_components is True:

        def dag_json():
            return merged_dag_content(ep_proto, components).dict()
    else:

        def dag_json():
            return component_dag_content(components).dict()
    return dag_json


def ishashable(x):
    """Is x hashable?

    Examples
    --------
    >>> ishashable(1)
    True
    >>> ishashable([1])
    False
    """
    try:
        hash(x)
        return True
    except TypeError:
        return False


def _execute_task(arg, cache):
    """Do the actual work of collecting data and executing a function.

    Examples
    --------
    >>> cache = {'x': 1, 'y': 2}  # Compute tasks against a cache
    >>> _execute_task((add, 'x', 1), cache)  # Compute task in naive manner
    2
    >>> _execute_task((add, (inc, 'x'), 1), cache)  # Support nested computation
    3
    >>> _execute_task('x', cache)  # Also grab data from cache
    1
    >>> list(_execute_task(['x', 'y'], cache))  # Support nested lists
    [1, 2]
    >>> list(map(list, _execute_task([['x', 'y'], ['y', 'x']], cache)))
    [[1, 2], [2, 1]]
    >>> _execute_task('foo', cache)  # Passes through on non-keys
    'foo'
    """
    if isinstance(arg, list):
        return [_execute_task(a, cache) for a in arg]
    if istask(arg):
        func, args = arg[0], arg[1:]
        return func(*(_execute_task(a, cache) for a in args))
    if not ishashable(arg):
        return arg
    if arg in cache:
        return cache[arg]
    return arg


def lists_to_tuples(res, keys):
    if isinstance(keys, list):
        return tuple(lists_to_tuples(r, k) for r, k in zip(res, keys))
    return res


def _toposort(dsk, keys=None, returncycle=False, dependencies=None):
    """Stack-based depth-first search traversal.

    This is based on Tarjan's method for topological sorting (see wikipedia for pseudocode).
    """
    if keys is None:
        keys = dsk
    elif not isinstance(keys, list):
        keys = [keys]
    if not returncycle:
        ordered = []
    completed = set()
    seen = set()
    if dependencies is None:
        dependencies = {k: get_dependencies(dsk, k) for k in dsk}
    for key in keys:
        if key in completed:
            continue
        nodes = [key]
        while nodes:
            cur = nodes[-1]
            if cur in completed:
                nodes.pop()
                continue
            seen.add(cur)
            next_nodes = []
            for nxt in dependencies[cur]:
                if nxt not in completed:
                    if nxt in seen:
                        cycle = [nxt]
                        while nodes[-1] != nxt:
                            cycle.append(nodes.pop())
                        cycle.append(nodes.pop())
                        cycle.reverse()
                        if returncycle:
                            return cycle
                        cycle = '->'.join(str(x) for x in cycle)
                        raise RuntimeError('Cycle detected in task graph: %s' % cycle)
                    next_nodes.append(nxt)
            if next_nodes:
                nodes.extend(next_nodes)
            else:
                if not returncycle:
                    ordered.append(cur)
                completed.add(cur)
                seen.remove(cur)
                nodes.pop()
    if returncycle:
        return []
    return ordered


def toposort(dsk, dependencies=None):
    """Return a list of keys of task graph sorted in topological order."""
    return _toposort(dsk, dependencies=dependencies)


def get(dsk: dict, out: Sequence[str], cache: dict=None, sortkeys: List[str]=None):
    """Get value from the task graphs.

    Parameters
    ----------
    dsk
        task graph dict
    out
        sequence of output keys which should be retrieved as results of running
        `get()` over the `dsk`.
    cache
        cache dict for fast in-memory lookups of previously computed results.
    sortkeys
        topologically sorted keys

    Examples
    --------
    >>> d = {'x': 1, 'y': (inc, 'x')}
    >>> get(d, 'x')
    1
    >>> get(d, 'y')
    2
    >>> get(d, 'y', sortkeys=['x', 'y'])
    2
    """
    for k in (flatten(out) if isinstance(out, list) else [out]):
        if k not in dsk:
            raise KeyError(f'{k} is not a key in the graph')
    if cache is None:
        cache = {}
    if sortkeys is None:
        sortkeys = toposort(dsk)
    for key in sortkeys:
        task = dsk[key]
        result = _execute_task(task, cache)
        cache[key] = result
    result = _execute_task(out, cache)
    if isinstance(out, list):
        result = lists_to_tuples(result, out)
    return result


class Token:
    """A token object.

    Used to express certain objects in the traversal of a task or pattern.
    """

    def __init__(self, name):
        self.name = name

    def __repr__(self):
        return self.name


END = Token('end')


def isdir(path: Any) ->bool:
    try:
        return os.path.isdir(path)
    except TypeError:
        return False


class SequencesDataset(Dataset):
    meta: dict
    split: str
    dataset_path = str
    label_to_names = dict
    num_classes: int
    path_list: list

    def __init__(self, data, cache_dir='./logs/cache', use_cache=False, num_points=65536, ignored_label_inds=[0], predicting=False, **kwargs):
        super().__init__()
        self.name = 'Dataset'
        self.ignored_label_inds = ignored_label_inds
        kwargs['cache_dir'] = cache_dir
        kwargs['use_cache'] = use_cache
        kwargs['num_points'] = num_points
        kwargs['ignored_label_inds'] = ignored_label_inds
        self.cfg = Config(kwargs)
        self.predicting = predicting
        if not predicting:
            self.on_fit(data)
        else:
            self.on_predict(data)

    @property
    def color_map(self):
        return self.meta['color_map']

    def on_fit(self, dataset_path):
        self.split = basename(dataset_path)
        self.load_meta(dirname(dataset_path))
        self.dataset_path = dataset_path
        self.label_to_names = self.get_label_to_names()
        self.num_classes = len(self.label_to_names) - len(self.ignored_label_inds)
        self.make_datasets()

    def load_meta(self, root_dir):
        meta_file = join(root_dir, 'meta.yaml')
        if not exists(meta_file):
            raise ValueError(f'The {root_dir} should contain a `meta.yaml` file about the pointcloud sequences.')
        with open(meta_file) as f:
            self.meta = yaml.safe_load(f)
        self.label_to_names = self.get_label_to_names()
        self.num_classes = len(self.label_to_names)
        with open(meta_file) as f:
            self.meta = yaml.safe_load(f)
        remap_dict_val = self.meta['learning_map']
        max_key = max(remap_dict_val.keys())
        remap_lut_val = np.zeros(max_key + 100, dtype=np.int32)
        remap_lut_val[list(remap_dict_val.keys())] = list(remap_dict_val.values())
        self.remap_lut_val = remap_lut_val

    def make_datasets(self):
        self.path_list = []
        for seq in os.listdir(self.dataset_path):
            sequence_path = join(self.dataset_path, seq)
            directories = [f for f in os.listdir(sequence_path) if isdir(join(sequence_path, f)) and f != 'labels']
            assert len(directories) == 1
            scan_dir = join(sequence_path, directories[0])
            for scan_name in os.listdir(scan_dir):
                self.path_list.append(join(scan_dir, scan_name))

    def on_predict(self, data):
        if isinstance(data, list):
            if not all(isfile(p) for p in data):
                raise ValueError('The predict input data takes only a list of paths or a directory.')
            root_dir = split(data[0])[0]
        elif isinstance(data, str):
            if not isdir(data) and not isfile(data):
                raise ValueError('The predict input data takes only a list of paths or a directory.')
            if isdir(data):
                root_dir = data
                data = [os.path.join(root_dir, f) for f in os.listdir(root_dir) if '.bin' in f]
            elif isfile(data):
                root_dir = dirname(data)
                data = [data]
            else:
                raise ValueError('The predict input data takes only a list of paths or a directory.')
        else:
            raise ValueError('The predict input data takes only a list of paths or a directory.')
        self.path_list = data
        self.split = 'predict'
        self.load_meta(root_dir)

    def get_label_to_names(self):
        """Returns a label to names dictionary object.

        Returns:
            A dict where keys are label numbers and
            values are the corresponding names.
        """
        return self.meta['label_to_names']

    def __getitem__(self, index):
        data = self.get_data(index)
        data['attr'] = self.get_attr(index)
        return data

    def get_data(self, idx):
        pc_path = self.path_list[idx]
        points = DataProcessing.load_pc_kitti(pc_path)
        folder, file = split(pc_path)
        if self.predicting:
            label_path = join(folder, file[:-4] + '.label')
        else:
            label_path = join(folder, '../labels', file[:-4] + '.label')
        if not exists(label_path):
            labels = np.zeros(np.shape(points)[0], dtype=np.int32)
            if self.split not in ['test', 'all']:
                raise FileNotFoundError(f' Label file {label_path} not found')
        else:
            labels = DataProcessing.load_label_kitti(label_path, self.remap_lut_val).astype(np.int32)
        data = {'point': points[:, 0:3], 'feat': None, 'label': labels}
        return data

    def get_attr(self, idx):
        pc_path = self.path_list[idx]
        folder, file = split(pc_path)
        _, seq = split(split(folder)[0])
        name = f'{seq}_{file[:-4]}'
        pc_path = str(pc_path)
        attr = {'idx': idx, 'name': name, 'path': pc_path, 'split': self.split}
        return attr

    def __len__(self):
        return len(self.path_list)

    def get_split(self, *_):
        return self


@functools.lru_cache()
def _raise_beta_warning(message: str, stacklevel: int=6):
    rank_zero_warn(f'{message} The API and functionality may change without warning in future releases. More details: https://lightning-flash.readthedocs.io/en/latest/stability.html', stacklevel=stacklevel, category=UserWarning)


def beta(message: str='This feature is currently in Beta.'):
    """The beta decorator is used to indicate that a particular feature is in Beta. A callable or type that has
    been marked as beta will give a ``UserWarning`` when it is called or instantiated. This designation should be
    used following the description given in :ref:`stability`.

    Args:
        message: The message to include in the warning.

    Examples
    ________

    .. testsetup::

        >>> import pytest

    .. doctest::

        >>> from flash.core.utilities.stability import beta
        >>> @beta()
        ... class MyBetaFeature:
        ...     pass
        ...
        >>> with pytest.warns(UserWarning, match="This feature is currently in Beta."):
        ...     MyBetaFeature()  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
        ...
        <...>
        >>> @beta("This feature is currently in Beta with a custom message.")
        ... class MyBetaFeatureWithCustomMessage:
        ...     pass
        ...
        >>> with pytest.warns(UserWarning, match="This feature is currently in Beta with a custom message."):
        ...     MyBetaFeatureWithCustomMessage()  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
        ...
        <...>
    """

    def decorator(callable: Union[Callable, Type]):
        if inspect.isclass(callable):
            callable.__init__ = decorator(callable.__init__)
            return callable

        @functools.wraps(callable)
        def wrapper(*args, **kwargs):
            _raise_beta_warning(message)
            return callable(*args, **kwargs)
        return wrapper
    return decorator


class Traverser:
    """Traverser interface for tasks.

    Class for storing the state while performing a preorder-traversal of a
    task.

    Parameters
    ----------
    term : task
        The task to be traversed

    Attributes
    ----------
    term
        The current element in the traversal
    current
        The head of the current element in the traversal. This is simply `head`
        applied to the attribute `term`.
    """

    def __init__(self, term, stack=None):
        self.term = term
        if not stack:
            self._stack = deque([END])
        else:
            self._stack = stack

    def __iter__(self):
        while self.current is not END:
            yield self.current
            self.next()

    def copy(self):
        """Copy the traverser in its current state.

        This allows the traversal to be pushed onto a stack, for easy backtracking.
        """
        return Traverser(self.term, deque(self._stack))

    def next(self):
        """Proceed to the next term in the preorder traversal."""
        subterms = args(self.term)
        if not subterms:
            self.term = self._stack.pop()
        else:
            self.term = subterms[0]
            self._stack.extend(reversed(subterms[1:]))

    @property
    def current(self):
        return head(self.term)

    def skip(self):
        """Skip over all subterms of the current level in the traversal."""
        self.term = self._stack.pop()


def subs(task, key, val):
    """Perform a substitution on a task.

    Examples
    --------
    >>> subs((inc, 'x'), 'x', 1)  # doctest: +SKIP
    (inc, 1)
    """
    type_task = type(task)
    if not (type_task is tuple and task and callable(task[0])):
        try:
            if type_task is type(key) and task == key:
                return val
        except Exception:
            pass
        if type_task is list:
            return [subs(x, key, val) for x in task]
        return task
    newargs = []
    for arg in task[1:]:
        type_arg = type(arg)
        if type_arg is tuple and arg and callable(arg[0]):
            arg = subs(arg, key, val)
        elif type_arg is list:
            arg = [subs(x, key, val) for x in arg]
        elif type_arg is type(key):
            try:
                if len(arg) == len(key) and all(type(aa) is type(bb) and aa == bb for aa, bb in zip(arg, key)):
                    arg = val
            except (TypeError, AttributeError):
                if arg == key:
                    arg = val
        newargs.append(arg)
    return task[:1] + tuple(newargs)


class RewriteRule:
    """A rewrite rule.

    Expresses `lhs` -> `rhs`, for variables `vars`.

    Parameters
    ----------
    lhs : task
        The left-hand-side of the rewrite rule.
    rhs : task or function
        The right-hand-side of the rewrite rule. If it's a task, variables in
        `rhs` will be replaced by terms in the subject that match the variables
        in `lhs`. If it's a function, the function will be called with a dict
        of such matches.
    vars: tuple, optional
        Tuple of variables found in the lhs. Variables can be represented as
        any hashable object; a good convention is to use strings. If there are
        no variables, this can be omitted.

    Examples
    --------
    Here's a `RewriteRule` to replace all nested calls to `list`, so that
    `(list, (list, 'x'))` is replaced with `(list, 'x')`, where `'x'` is a
    variable.

    >>> lhs = (list, (list, 'x'))
    >>> rhs = (list, 'x')
    >>> variables = ('x',)
    >>> rule = RewriteRule(lhs, rhs, variables)

    Here's a more complicated rule that uses a callable right-hand-side. A
    callable `rhs` takes in a dictionary mapping variables to their matching
    values. This rule replaces all occurrences of `(list, 'x')` with `'x'` if
    `'x'` is a list itself.

    >>> lhs = (list, 'x')
    >>> def repl_list(sd):
    ...     x = sd['x']
    ...     if isinstance(x, list):
    ...         return x
    ...     else:
    ...         return list, x
    >>> rule = RewriteRule(lhs, repl_list, variables)
    """

    def __init__(self, lhs, rhs, vars=()):
        if not isinstance(vars, tuple):
            raise TypeError('vars must be a tuple of variables')
        self.lhs = lhs
        if callable(rhs):
            self.subs = rhs
        else:
            self.subs = self._apply
        self.rhs = rhs
        self._varlist = [t for t in Traverser(lhs) if t in vars]
        self.vars = tuple(sorted(set(self._varlist)))

    def _apply(self, sub_dict):
        term = self.rhs
        for key, val in sub_dict.items():
            term = subs(term, key, val)
        return term

    def __str__(self):
        return f'RewriteRule({self.lhs}, {self.rhs}, {self.vars})'

    def __repr__(self):
        return str(self)


class Node(tuple):
    """A Discrimination Net node."""
    __slots__ = ()

    def __new__(cls, edges=None, patterns=None):
        edges = edges if edges else {}
        patterns = patterns if patterns else []
        return tuple.__new__(cls, (edges, patterns))

    @property
    def edges(self):
        """A dictionary, where the keys are edges, and the values are nodes."""
        return self[0]

    @property
    def patterns(self):
        """A list of all patterns that currently match at this node."""
        return self[1]


VAR = Token('?')


def _match(S, N):
    """Structural matching of term S to discrimination net node N."""
    stack = deque()
    restore_state_flag = False
    matches = ()
    while True:
        if S.current is END:
            yield N.patterns, matches
        try:
            n = N.edges.get(S.current, None)
            if n and not restore_state_flag:
                stack.append((S.copy(), N, matches))
                N = n
                S.next()
                continue
        except TypeError:
            pass
        n = N.edges.get(VAR, None)
        if n:
            restore_state_flag = False
            matches = matches + (S.term,)
            S.skip()
            N = n
            continue
        try:
            S, N, matches = stack.pop()
            restore_state_flag = True
        except Exception:
            return


def _process_match(rule, syms):
    """Process a match to determine if it is correct, and to find the correct substitution that will convert the
    term into the pattern.

    Parameters
    ----------
    rule : RewriteRule
    syms : iterable
        Iterable of subterms that match a corresponding variable.

    Returns
    -------
    A dictionary of {vars : subterms} describing the substitution to make the
    pattern equivalent with the term. Returns `None` if the match is
    invalid.
    """
    subs = {}
    varlist = rule._varlist
    if not len(varlist) == len(syms):
        raise RuntimeError("length of varlist doesn't match length of syms.")
    for v, s in zip(varlist, syms):
        if v in subs and subs[v] != s:
            return None
        subs[v] = s
    return subs


def _bottom_up(net, term):
    if istask(term):
        term = (head(term),) + tuple(_bottom_up(net, t) for t in args(term))
    elif isinstance(term, list):
        term = [_bottom_up(net, t) for t in args(term)]
    return net._rewrite(term)


def _top_level(net, term):
    return net._rewrite(term)


strategies = {'top_level': _top_level, 'bottom_up': _bottom_up}


class RuleSet:
    """A set of rewrite rules.

    Forms a structure for fast rewriting over a set of rewrite rules. This
    allows for syntactic matching of terms to patterns for many patterns at
    the same time.

    Examples
    --------

    >>> def f(*args): pass
    >>> def g(*args): pass
    >>> def h(*args): pass
    >>> from operator import add

    >>> rs = RuleSet(                 # Make RuleSet with two Rules
    ...         RewriteRule((add, 'x', 0), 'x', ('x',)),
    ...         RewriteRule((f, (g, 'x'), 'y'),
    ...                     (h, 'x', 'y'),
    ...                     ('x', 'y')))

    >>> rs.rewrite((add, 2, 0))       # Apply ruleset to single task
    2

    >>> rs.rewrite((f, (g, 'a', 3)))  # doctest: +SKIP
    (h, 'a', 3)

    >>> dsk = {'a': (add, 2, 0),      # Apply ruleset to full dask graph
    ...        'b': (f, (g, 'a', 3))}

    Attributes
    ----------
    rules : list
        A list of `RewriteRule`s included in the `RuleSet`.
    """

    def __init__(self, *rules):
        """Create a `RuleSet` for a number of rules.

        Parameters
        ----------
        rules
            One or more instances of RewriteRule
        """
        self._net = Node()
        self.rules = []
        for p in rules:
            self.add(p)

    def add(self, rule):
        """Add a rule to the RuleSet.

        Parameters
        ----------
        rule : RewriteRule
        """
        if not isinstance(rule, RewriteRule):
            raise TypeError('rule must be instance of RewriteRule')
        list_vars = rule.vars
        curr_node = self._net
        ind = len(self.rules)
        for t in Traverser(rule.lhs):
            prev_node = curr_node
            if t in list_vars:
                t = VAR
            if t in curr_node.edges:
                curr_node = curr_node.edges[t]
            else:
                curr_node.edges[t] = Node()
                curr_node = curr_node.edges[t]
        prev_node.edges[t].patterns.append(ind)
        self.rules.append(rule)

    def iter_matches(self, term):
        """A generator that lazily finds matchings for term from the RuleSet.

        Parameters
        ----------
        term : task

        Yields
        ------
        Tuples of `(rule, subs)`, where `rule` is the rewrite rule being
        matched, and `subs` is a dictionary mapping the variables in the lhs
        of the rule to their matching values in the term.
        """
        S = Traverser(term)
        for m, syms in _match(S, self._net):
            for i in m:
                rule = self.rules[i]
                subs = _process_match(rule, syms)
                if subs is not None:
                    yield rule, subs

    def _rewrite(self, term):
        """Apply the rewrite rules in RuleSet to top level of term."""
        for rule, sd in self.iter_matches(term):
            term = rule.subs(sd)
            break
        return term

    def rewrite(self, task, strategy='bottom_up'):
        """Apply the `RuleSet` to `task`.

        This applies the most specific matching rule in the RuleSet to the
        task, using the provided strategy.

        Parameters
        ----------
        task: a task
            The task to be rewritten
        strategy: str, optional
            The rewriting strategy to use. Options are "bottom_up" (default),
            or "top_level".

        Examples
        --------
        Suppose there was a function `add` that returned the sum of 2 numbers,
        and another function `double` that returned twice its input:

        >>> add = lambda x, y: x + y
        >>> double = lambda x: 2*x

        Now suppose `double` was *significantly* faster than `add`, so
        you'd like to replace all expressions `(add, x, x)` with `(double,
        x)`, where `x` is a variable. This can be expressed as a rewrite rule:

        >>> rule = RewriteRule((add, 'x', 'x'), (double, 'x'), ('x',))
        >>> rs = RuleSet(rule)

        This can then be applied to terms to perform the rewriting:

        >>> term = (add, (add, 2, 2), (add, 2, 2))
        >>> rs.rewrite(term)  # doctest: +SKIP
        (double, (double, 2))

        If we only wanted to apply this to the top level of the term, the
        `strategy` kwarg can be set to "top_level".

        >>> rs.rewrite(term)  # doctest: +SKIP
        (double, (add, 2, 2))
        """
        return strategies[strategy](self, task)


def getcycle(d, keys):
    """Return a list of nodes that form a cycle if graph is not a DAG. Returns an empty list if no cycle is found.
    ``keys`` may be a single key or list of keys.

    Examples
    --------
    >>> d = {'x': (inc, 'z'), 'y': (inc, 'x'), 'z': (inc, 'y')}
    >>> getcycle(d, 'x')
    ['x', 'z', 'y', 'x']

    See Also
    --------
    isdag
    """
    return _toposort(d, keys=keys, returncycle=True)


def isdag(d, keys):
    """Does graph form a directed acyclic graph when calculating keys? ``keys`` may be a single key or list of
    keys.

    Examples
    --------
    >>> isdag({'x': 0, 'y': (inc, 'x')}, 'y')
    True
    >>> isdag({'x': (inc, 'y'), 'y': (inc, 'x')}, 'y')
    False

    See Also
    --------
    getcycle
    """
    return not getcycle(d, keys)


def _verify_no_cycles(dsk: Dict[str, tuple], out_keys: List[str], endpoint_name: str):
    if not isdag(dsk, keys=out_keys):
        cycle = getcycle(dsk, keys=out_keys)
        raise RuntimeError(f'Cycle detected when attepting to build DAG for endpoint: `{endpoint_name}`. This cycle is formed by connections between the following nodes: {cycle}')


def cull(dsk, keys):
    """Return new task graph with only the tasks required to calculate keys.

    In other words, remove unnecessary tasks from task graph.
    ``keys`` may be a single key or list of keys.

    Examples
    --------
    >>> d = {'x': 1, 'y': (inc, 'x'), 'out': (add, 'x', 10)}
    >>> dsk, dependencies = cull(d, 'out')  # doctest: +SKIP
    >>> dsk  # doctest: +SKIP
    {'x': 1, 'out': (add, 'x', 10)}
    >>> dependencies  # doctest: +SKIP
    {'x': set(), 'out': set(['x'])}

    Returns
    -------
    dsk: culled graph
    dependencies: Dict mapping {key: [deps]}.  Useful side effect to accelerate
        other optimizations, notably fuse.
    """
    if not isinstance(keys, (list, set)):
        keys = [keys]
    seen = set()
    dependencies = dict()
    out = {}
    work = list(set(flatten(keys)))
    while work:
        new_work = []
        for k in work:
            dependencies_k = get_dependencies(dsk, k, as_list=True)
            out[k] = dsk[k]
            dependencies[k] = dependencies_k
            for d in dependencies_k:
                if d not in seen:
                    seen.add(d)
                    new_work.append(d)
        work = new_work
    return out, dependencies


def _flat_set(x):
    if x is None:
        return set()
    if isinstance(x, set):
        return x
    if not isinstance(x, (list, set)):
        x = [x]
    return set(x)


def inline(dsk, keys=None, inline_constants=True, dependencies=None):
    """Return new dask with the given keys inlined with their values.

    Inlines all constants if ``inline_constants`` keyword is True. Note that
    the constant keys will remain in the graph, to remove them follow
    ``inline`` with ``cull``.

    Examples
    --------
    >>> d = {'x': 1, 'y': (inc, 'x'), 'z': (add, 'x', 'y')}
    >>> inline(d)  # doctest: +SKIP
    {'x': 1, 'y': (inc, 1), 'z': (add, 1, 'y')}
    >>> inline(d, keys='y')  # doctest: +SKIP
    {'x': 1, 'y': (inc, 1), 'z': (add, 1, (inc, 1))}
    >>> inline(d, keys='y', inline_constants=False)  # doctest: +SKIP
    {'x': 1, 'y': (inc, 1), 'z': (add, 'x', (inc, 'x'))}
    """
    if dependencies and isinstance(next(iter(dependencies.values())), list):
        dependencies = {k: set(v) for k, v in dependencies.items()}
    keys = _flat_set(keys)
    if dependencies is None:
        dependencies = {k: get_dependencies(dsk, k) for k in dsk}
    if inline_constants:
        keys.update(k for k, v in dsk.items() if ishashable(v) and v in dsk or not dependencies[k] and not istask(v))
    replaceorder = toposort({k: dsk[k] for k in keys if k in dsk}, dependencies=dependencies)
    keysubs = {}
    for key in replaceorder:
        val = dsk[key]
        for dep in (keys & dependencies[key]):
            replace = keysubs.get(dep, dsk[dep])
            val = subs(val, dep, replace)
        keysubs[key] = val
    dsk2 = keysubs.copy()
    for key, val in dsk.items():
        if key not in dsk2:
            for item in (keys & dependencies[key]):
                val = subs(val, item, keysubs[item])
            dsk2[key] = val
    return dsk2


def inline_functions(dsk, output, fast_functions=None, inline_constants=False, dependencies=None):
    """Inline cheap functions into larger operations.

    Examples
    --------
    >>> double = lambda x: x*2  # doctest: +SKIP
    >>> dsk = {'out': (add, 'i', 'd'),  # doctest: +SKIP
    ...        'i': (inc, 'x'),
    ...        'd': (double, 'y'),
    ...        'x': 1, 'y': 1}
    >>> inline_functions(dsk, [], [inc])  # doctest: +SKIP
    {'out': (add, (inc, 'x'), 'd'),
     'd': (double, 'y'),
     'x': 1, 'y': 1}

    Protect output keys.  In the example below ``i`` is not inlined because it
    is marked as an output key.

    >>> inline_functions(dsk, ['i', 'out'], [inc, double])  # doctest: +SKIP
    {'out': (add, 'i', (double, 'y')),
     'i': (inc, 'x'),
     'x': 1, 'y': 1}
    """
    if not fast_functions:
        return dsk
    output = set(output)
    fast_functions = set(fast_functions)
    if dependencies is None:
        dependencies = {k: get_dependencies(dsk, k) for k in dsk}
    dependents = reverse_dict(dependencies)

    def inlinable(v):
        try:
            return functions_of(v).issubset(fast_functions)
        except TypeError:
            return False
    keys = [k for k, v in dsk.items() if istask(v) and dependents[k] and k not in output and inlinable(v)]
    if keys:
        dsk = inline(dsk, keys, inline_constants=inline_constants, dependencies=dependencies)
        for k in keys:
            del dsk[k]
    return dsk


def build_composition(endpoint_protocol: 'EndpointProtocol', components: Dict[str, 'ModelComponent'], connections: List['Connection']) ->'TaskComposition':
    """Build a composed graph.

    Notes on easy sources to introduce bugs.

    ::

            Input Data
        --------------------
            a  b  c   d
            |  |  |   | \\\\
             \\ | / \\  |  ||
              C_2   C_1  ||
            /  |     | \\ //
           /   |    /   *
        RES_2  |   |   // \\
               |   |  //   RES_1
                \\  | //
                C_2_1
                  |
                RES_3
        ---------------------
              Output Data

    Because there are connections between ``C_1 -> C_2_1`` and
    ``C_2 -> C_2_1`` we can eliminate the ``serialize <-> deserialize``
    tasks for the data transfered between these components. We need to be
    careful to not eliminate the ``serialize`` or ``deserialize`` tasks
    entirely though. In the case shown above, it is apparent ``RES_1`` &
    ``RES_2``. still need the ``serialize`` function, but the same also applies
    for ``deserialize``. Consider the example below with the same composition &
    connections as above:

    ::
            Input Data
        --------------------
            a  b  c   d
            |  |  |   | \\\\
             \\ | /| \\ |  \\\\
              C_2 |  C_1  ||
            /  |  |   @\\  ||
           /   |  |   @ \\ //
        RES_2  |  |  @   *
               |  | @  // \\
                \\ | @ //   RES_1
                 C_2_1
                  |
                RES_3
        ---------------------
              Output Data

    Though we are using the same composition, the endpoints have been changed so
    that the previous result of ``C_1``-> ``C_2_1`` is now being provided by
    input ``c``. However, there is still a connection between ``C_1`` and
    ``C_2_1`` which is denoted by the ``@`` symbols... Though the first
    example (shown at the top of this docstring) would be able to eliminate
    ``C_2_1 deserailize``from ``C_2`` / ``C_1``, we see here that since
    endpoints define the path through the DAG, we cannot eliminate them
    entirely either.
    """
    initial_task_dsk = _process_initial(endpoint_protocol, components)
    dsk_tgt_src_connections = {}
    for connection in connections:
        source_dsk = f'{connection.source_component}.outputs.{connection.source_key}'
        target_dsk = f'{connection.target_component}.inputs.{connection.target_key}'
        dsk_tgt_src_connections[target_dsk] = identity, source_dsk
    rewrite_ruleset = RuleSet()
    for dsk_payload_target_serial in initial_task_dsk.payload_tasks_dsk.keys():
        dsk_payload_target, _serial_ident = dsk_payload_target_serial.rsplit('.', maxsplit=1)
        if _serial_ident != 'serial':
            raise RuntimeError(f'dsk_payload_target_serial={dsk_payload_target_serial}, dsk_payload_target={dsk_payload_target}, _serial_ident={_serial_ident}')
        if dsk_payload_target in dsk_tgt_src_connections:
            lhs = dsk_tgt_src_connections[dsk_payload_target]
            rhs = initial_task_dsk.merged_dsk[dsk_payload_target]
            rule = RewriteRule(lhs, rhs, vars=())
            rewrite_ruleset.add(rule)
    io_subgraphs_merged = merge(initial_task_dsk.merged_dsk, dsk_tgt_src_connections, initial_task_dsk.result_tasks_dsk, initial_task_dsk.payload_tasks_dsk)
    rewritten_dsk = valmap(rewrite_ruleset.rewrite, io_subgraphs_merged)
    culled_dsk, culled_deps = cull(rewritten_dsk, initial_task_dsk.output_keys)
    _verify_no_cycles(culled_dsk, initial_task_dsk.output_keys, endpoint_protocol.name)
    inlined = inline_functions(culled_dsk, initial_task_dsk.output_keys, fast_functions=[identity], inline_constants=True, dependencies=culled_deps)
    inlined_culled_dsk, inlined_culled_deps = cull(inlined, initial_task_dsk.output_keys)
    _verify_no_cycles(inlined_culled_dsk, initial_task_dsk.output_keys, endpoint_protocol.name)
    toposort_keys = toposort(inlined_culled_dsk)
    res = TaskComposition(dsk=inlined_culled_dsk, sortkeys=toposort_keys, get_keys=initial_task_dsk.output_keys, ep_dsk_input_keys=initial_task_dsk.payload_dsk_map, ep_dsk_output_keys=initial_task_dsk.result_dsk_map, pre_optimization_dsk=initial_task_dsk.merged_dsk)
    return res


def setup_http_app(composition: 'Composition', debug: bool) ->'FastAPI':
    app = FastAPI(debug=debug, version=__version__, title='FlashServe')
    app.get('/flashserve/alive', name='alive', description='If you can reach this endpoint, the server is runnning.', response_model=Alive)(_build_alive_check())
    _no_optimization_dsk = build_composition(endpoint_protocol=first(composition.endpoint_protocols.values()), components=composition.components, connections=composition.connections)
    pth = Path(__file__).parent.joinpath('templates')
    templates = Jinja2Templates(directory=str(pth.absolute()))
    app.get('/flashserve/component_dags', name='component_dags', summary='HTML Rendering of Component DAGs', response_class=HTMLResponse)(_build_visualization(dsk_composition=_no_optimization_dsk, templates=templates, no_optimization=True))
    app.get('/flashserve/dag_json', name='components JSON DAG', summary='JSON representation of component DAG', response_model=ComponentJSON)(_build_dag_json(components=composition.components, ep_proto=None, show_connected_components=False))
    for ep_name, ep_proto in composition.endpoint_protocols.items():
        dsk = build_composition(endpoint_protocol=ep_proto, components=composition.components, connections=composition.connections)
        RequestModel = ep_proto.request_model
        ResponseModel = ep_proto.response_model
        app.post(f'{ep_proto.route}', name=ep_name, tags=[ep_name], summary='Perform a Compution.', description='Computes results of DAG defined by these components & endpoint.', response_model=ResponseModel)(_build_endpoint(RequestModel, dsk, ResponseModel))
        app.get(f'{ep_proto.route}/meta', name=f'{ep_name} meta schema', tags=[ep_name], summary='OpenAPI schema', description="OpenAPI schema for this endpoints's compute route.")(_build_meta(RequestModel))
        app.get(f'{ep_proto.route}/dag', name=f'{ep_name} DAG Visualization', tags=[ep_name], summary='HTML Rendering of DAG', description='Displays an html image rendering the DAG of functions & components executed to reach the endpoint outputs.', response_class=HTMLResponse)(_build_visualization(dsk, templates))
        app.get(f'{ep_proto.route}/dag_json', name=f'{ep_name} JSON DAG', tags=[ep_name], summary='JSON representatino of DAG', response_model=MergedJSON)(_build_dag_json(components=composition.components, ep_proto=ep_proto, show_connected_components=True))
    return app


class ServerMixin:
    """Start a server to serve a composition.

    debug     If the server should be started up in debug mode. By default, False. testing     If the server should
    return the ``app`` instance instead of blocking     the process (via running the ``app`` in ``uvicorn``). This is
    used     when taking advantage of a server ``TestClient``. By default, False
    """
    DEBUG: bool
    TESTING: bool

    def http_app(self) ->'FastAPI':
        return setup_http_app(composition=self, debug=self.DEBUG)

    def serve(self, host: str='127.0.0.1', port: int=8000):
        """Start a server to serve a composition.

        Parameters
        ----------
        host
            host address to run the server on
        port
            port number to expose the running server on
        """
        if FLASH_DISABLE_SERVE:
            return
        if not self.TESTING:
            app = self.http_app()
            uvicorn.run(app, host=host, port=port)
        return self.http_app()


class FineTuningHooks:
    """Hooks to be used in Task and FlashBaseTuning."""

    def modules_to_freeze(self) ->Optional[Union[Module, Iterable[Union[Module, Iterable]]]]:
        """Return the name(s) of the module attributes of the model to be frozen."""
        return None


LOSS_FN_TYPE = TypeVar('LOSS_FN_TYPE', Callable, Mapping, Sequence, None)


LR_SCHEDULER_TYPE = TypeVar('LR_SCHEDULER_TYPE', str, Callable, Tuple[str, Dict[str, Any]], Tuple[str, Dict[str, Any], Dict[str, Any]], None)


MODEL_TYPE = TypeVar('MODEL_TYPE', nn.Module, None)


OPTIMIZER_TYPE = TypeVar('OPTIMIZER_TYPE', str, Callable, Tuple[str, Dict[str, Any]], None)


OUTPUT_TRANSFORM_TYPE = TypeVar('OUTPUT_TRANSFORM_TYPE', 'flash.core.data.io.output_transform.OutputTransform', None)


def get_callable_name(fn_or_class: Union[Callable, object]) ->str:
    return getattr(fn_or_class, '__name__', fn_or_class.__class__.__name__).lower()


def get_callable_dict(fn: Union[nn.Module, Callable, Mapping, Sequence]) ->Union[Dict, Mapping]:
    if isinstance(fn, nn.Module):
        return nn.ModuleDict({get_callable_name(fn): fn})
    if isinstance(fn, Mapping):
        return fn
    if isinstance(fn, Sequence):
        return {get_callable_name(f): f for f in fn}
    if callable(fn):
        return {get_callable_name(fn): fn}


vars = 'a', 'b', 'c'


def _defaults_from_env_vars(fn: Callable) ->Callable:
    """Copy of ``pytorch_lightning.trainer.connectors.env_vars_connector._defaults_from_env_vars``.

    Required to fix build error in readthedocs.
    """

    @wraps(fn)
    def insert_env_defaults(self, *args, **kwargs):
        cls = self.__class__
        if args:
            cls_arg_names = [arg[0] for arg in get_init_arguments_and_types(cls)]
            kwargs.update(dict(zip(cls_arg_names, args)))
        env_variables = vars(parse_env_variables(cls))
        kwargs = dict(list(env_variables.items()) + list(kwargs.items()))
        return fn(self, **kwargs)
    return insert_env_defaults


class ModuleWrapperBase:
    """The ``ModuleWrapperBase`` is a base for classes which wrap a ``LightningModule`` or an instance of
    ``ModuleWrapperBase``.

    This class ensures that trainer attributes are forwarded to any wrapped or nested
    ``LightningModule`` instances so that nested calls to ``.log`` are handled correctly. The ``ModuleWrapperBase`` is
    also stateful, meaning that a :class:`~flash.core.data.data_pipeline.DataPipelineState` can be attached. Attached
    state will be forwarded to any nested ``ModuleWrapperBase`` instances.
    """

    def __init__(self):
        super().__init__()
        self._children = []

    def __setattr__(self, key, value):
        if isinstance(value, (LightningModule, ModuleWrapperBase)):
            self._children.append(key)
        patched_attributes = ['_current_fx_name', '_current_hook_fx_name', '_results', '_data_pipeline_state']
        if isinstance(value, Trainer) or key in patched_attributes:
            if hasattr(self, '_children'):
                for child in self._children:
                    setattr(getattr(self, child), key, value)
        super().__setattr__(key, value)


class _ServeInputProcessor(nn.Module):

    def __init__(self, serve_input: 'ServeInput', collate_fn: Callable):
        super().__init__()
        self.serve_input = serve_input
        self.collate_fn = collate_fn

    def forward(self, sample: str):
        sample = self.serve_input._call_load_sample(sample)
        if not isinstance(sample, list):
            sample = [sample]
        sample = self.collate_fn(sample)
        return sample


class FuncModule(nn.Module):
    """This class is used to wrap a callable within a nn.Module and apply the wrapped function in `__call__`"""

    def __init__(self, func: Callable) ->None:
        super().__init__()
        self.func = func

    def forward(self, *args, **kwargs) ->Any:
        return self.func(*args, **kwargs)

    def __str__(self) ->str:
        return f'{self.__class__.__name__}({self.func.__name__})'

    def __repr__(self):
        return str(self.func)


def convert_to_modules(transforms: Optional[Dict[str, Callable]]):
    if transforms is None or isinstance(transforms, nn.Module):
        return transforms
    transforms = apply_to_collection(transforms, Callable, FuncModule, wrong_dtype=nn.Module)
    transforms = apply_to_collection(transforms, Mapping, nn.ModuleDict, wrong_dtype=nn.ModuleDict)
    transforms = apply_to_collection(transforms, Iterable, nn.ModuleList, wrong_dtype=(nn.ModuleList, nn.ModuleDict))
    return transforms


class ApplyToKeys(nn.Sequential):
    """The ``ApplyToKeys`` class is an ``nn.Sequential`` which applies the given transforms to the given keys from
    the input. When a single key is given, a single value will be passed to the transforms. When multiple keys are
    given, the corresponding values will be passed to the transforms as a list.

    Args:
        keys: The key (``str``) or sequence of keys (``Sequence[str]``) to extract and forward to the transforms.
        args: The transforms, passed to the ``nn.Sequential`` super constructor.
    """

    def __init__(self, keys: Union[str, Sequence[str]], *args):
        super().__init__(*(convert_to_modules(arg) for arg in args))
        if isinstance(keys, str):
            keys = [keys]
        self.keys = keys

    def forward(self, x: Mapping[str, Any]) ->Mapping[str, Any]:
        keys = list(filter(lambda key: key in x, self.keys))
        inputs = [x[key] for key in keys]
        result = {}
        result.update(x)
        if len(inputs) == 1:
            result[keys[0]] = super().forward(inputs[0])
        elif len(inputs) > 1:
            try:
                outputs = super().forward(inputs)
            except TypeError as e:
                raise Exception('Failed to apply transforms to multiple keys at the same time.') from e
            for i, key in enumerate(keys):
                result[key] = outputs[i]
        return result

    def __repr__(self):
        transform = list(self.children())
        keys = self.keys[0] if len(self.keys) == 1 else self.keys
        transform = transform[0] if len(transform) == 1 else transform
        return f'{self.__class__.__name__}(keys={repr(keys)}, transform={repr(transform)})'


def catch_url_error(fn):

    @functools.wraps(fn)
    def wrapper(*args, pretrained=False, **kwargs):
        try:
            return fn(*args, pretrained=pretrained, **kwargs)
        except urllib.error.URLError:
            if 'efficientdet' in kwargs.get('head', ''):
                kwargs['pretrained_backbone'] = False
            result = fn(*args, pretrained=False, **kwargs)
            rank_zero_warn('Failed to download pretrained weights for the selected backbone. The backbone has been created with `pretrained=False` instead. If you are loading from a local checkpoint, this warning can be safely ignored.', category=UserWarning)
            return result
    return wrapper


def from_icevision_detection(record: 'BaseRecord'):
    detection = record.detection
    result = {}
    if hasattr(detection, 'bboxes'):
        result['bboxes'] = [{'xmin': bbox.xmin, 'ymin': bbox.ymin, 'width': bbox.width, 'height': bbox.height} for bbox in detection.bboxes]
    masks = getattr(detection, 'masks', None)
    mask_array = getattr(detection, 'mask_array', None)
    if mask_array is not None or not _ICEVISION_GREATER_EQUAL_0_11_0:
        if not isinstance(mask_array, MaskArray) or len(mask_array.data) == 0:
            mask_array = MaskArray.from_masks(masks, record.height, record.width)
        result['masks'] = [mask.data[0] for mask in _split_mask_array(mask_array)]
    elif masks is not None:
        result['masks'] = masks
    if hasattr(detection, 'keypoints'):
        keypoints = detection.keypoints
        result['keypoints'] = []
        result['keypoints_metadata'] = []
        for keypoint in keypoints:
            keypoints_list = []
            for x, y, v in keypoint.xyv:
                keypoints_list.append({'x': x, 'y': y, 'visible': v})
            result['keypoints'].append(keypoints_list)
            result['keypoints_metadata'].append(keypoint.metadata)
    if getattr(detection, 'label_ids', None) is not None:
        result['labels'] = list(detection.label_ids)
    if getattr(detection, 'scores', None) is not None:
        result['scores'] = list(detection.scores)
    return result


def from_icevision_predictions(predictions: List['Prediction']):
    result = []
    for prediction in predictions:
        result.append(from_icevision_detection(prediction.pred))
    return result


def from_icevision_record(record: 'BaseRecord'):
    sample = {DataKeys.METADATA: {'size': getattr(record, 'original_size', (record.height, record.width)), 'output_size': (record.height, record.width)}}
    if getattr(record, 'record_id', None) is not None:
        sample[DataKeys.METADATA]['image_id'] = record.record_id
    if getattr(record, 'filepath', None) is not None:
        sample[DataKeys.METADATA]['filepath'] = record.filepath
    if record.img is not None:
        sample[DataKeys.INPUT] = record.img
        filepath = getattr(record, 'filepath', None)
        if filepath is not None:
            sample[DataKeys.METADATA]['filepath'] = filepath
    elif getattr(record, 'filepath', None) is not None:
        sample[DataKeys.INPUT] = record.filepath
    sample[DataKeys.TARGET] = from_icevision_detection(record)
    if getattr(record.detection, 'class_map', None) is not None:
        sample[DataKeys.METADATA]['class_map'] = record.detection.class_map
    return sample


tasks = ['flash.audio.classification', 'flash.audio.speech_recognition', 'flash.graph.classification', 'flash.image.classification', 'flash.image.detection', 'flash.image.face_detection', 'flash.image.instance_segmentation', 'flash.image.keypoint_detection', 'flash.image.segmentation', 'flash.image.style_transfer', 'flash.pointcloud.detection', 'flash.pointcloud.segmentation', 'flash.tabular.classification', 'flash.tabular.regression', 'flash.tabular.forecasting', 'flash.text.classification', 'flash.text.question_answering', 'flash.text.seq2seq.summarization', 'flash.text.seq2seq.translation', 'flash.video.classification']


def to_icevision_record(sample: Dict[str, Any]):
    record = BaseRecord([])
    metadata = sample.get(DataKeys.METADATA, None) or {}
    if 'image_id' in metadata:
        record_id_component = RecordIDRecordComponent()
        record_id_component.set_record_id(metadata['image_id'])
    component = ClassMapRecordComponent(tasks.detection)
    component.set_class_map(metadata.get('class_map', None))
    record.add_component(component)
    if isinstance(sample[DataKeys.INPUT], str):
        input_component = FilepathRecordComponent()
        input_component.set_filepath(sample[DataKeys.INPUT])
    else:
        if 'filepath' in metadata:
            input_component = FilepathRecordComponent()
            input_component.filepath = metadata['filepath']
        else:
            input_component = ImageRecordComponent()
        input_component.composite = record
        image = sample[DataKeys.INPUT]
        if isinstance(image, Tensor):
            image = image.permute(1, 2, 0).numpy()
        elif isinstance(image, Image.Image):
            image = np.array(image)
        input_component.set_img(image)
        record.add_component(OriginalSizeRecordComponent(metadata.get('size', image.shape[:2])))
    record.add_component(input_component)
    if DataKeys.TARGET in sample:
        if 'labels' in sample[DataKeys.TARGET]:
            labels_component = InstancesLabelsRecordComponent()
            labels_component.add_labels_by_id(sample[DataKeys.TARGET]['labels'])
            record.add_component(labels_component)
        if 'bboxes' in sample[DataKeys.TARGET]:
            bboxes = [BBox.from_xywh(bbox['xmin'], bbox['ymin'], bbox['width'], bbox['height']) for bbox in sample[DataKeys.TARGET]['bboxes']]
            bboxes_component = BBoxesRecordComponent()
            bboxes_component.set_bboxes(bboxes)
            record.add_component(bboxes_component)
        if _ICEVISION_GREATER_EQUAL_0_11_0:
            masks = sample[DataKeys.TARGET].get('masks', None)
            if masks is not None:
                component = InstanceMasksRecordComponent()
                if len(masks) > 0 and isinstance(masks[0], Mask):
                    component.set_masks(masks)
                else:
                    if len(masks) == 0 or not len(masks) == len(record.detection.bboxes) == len(record.detection.label_ids):
                        data = np.zeros((0, record.height, record.width), np.uint8)
                        labels_component.label_ids = []
                        bboxes_component.bboxes = []
                    else:
                        data = np.stack(masks, axis=0)
                    mask_array = MaskArray(data)
                    component.set_mask_array(mask_array)
                    component.set_masks(_split_mask_array(mask_array))
                record.add_component(component)
        else:
            mask_array = sample[DataKeys.TARGET].get('mask_array', None)
            if mask_array is not None:
                component = MasksRecordComponent()
                component.set_masks(mask_array)
                record.add_component(component)
        if 'keypoints' in sample[DataKeys.TARGET]:
            keypoints = []
            keypoints_metadata = sample[DataKeys.TARGET].get('keypoints_metadata', [None] * len(sample[DataKeys.TARGET]['keypoints']))
            for keypoints_list, keypoints_metadata in zip(sample[DataKeys.TARGET]['keypoints'], keypoints_metadata):
                xyv = []
                for keypoint in keypoints_list:
                    xyv.extend((keypoint['x'], keypoint['y'], keypoint['visible']))
                keypoints.append(KeyPoints.from_xyv(xyv, keypoints_metadata))
            component = KeyPointsRecordComponent()
            component.set_keypoints(keypoints)
            record.add_component(component)
    return record


def _effdet_validation_step(validation_step, batch, batch_idx):
    images = batch[0][0]
    batch[0][1]['img_scale'] = torch.ones_like(images[:, 0, 0, 0]).unsqueeze(1)
    batch[0][1]['img_size'] = (torch.ones_like(images[:, 0, 0, 0]) * images[0].shape[-1]).unsqueeze(1).repeat(1, 2)
    return validation_step(batch, batch_idx)


def _log_with_name_and_prog_bar_override(log, adapter, name, value, **kwargs):
    if 'prog_bar' not in kwargs:
        kwargs['prog_bar'] = True
    metric = name.split('/')[-1]
    metric = f'{_STAGES_PREFIX[adapter.trainer.state.stage]}_{metric}'
    return log(metric, value, **kwargs)


def wrap_icevision_adapter(adapter):
    if not isinstance(adapter.log, partial):
        adapter.log = partial(_log_with_name_and_prog_bar_override, adapter.log, adapter)
    if isinstance(adapter, EffDetModelAdapter) and not isinstance(adapter.validation_step, partial):
        adapter.validation_step = partial(_effdet_validation_step, adapter.validation_step)
    return adapter


class A:

    def __call__(self, x):
        return True


class IceVisionTransformAdapter(nn.Module):
    """
    Args:
        transform: list of transformation functions to apply

    """

    def __init__(self, transform: List[Callable]):
        super().__init__()
        self.transform = A.Adapter(transform)

    def forward(self, x):
        record = to_icevision_record(x)
        record = self.transform(record)
        return from_icevision_record(record)


class DefaultGraphHead(nn.Module):

    def __init__(self, hidden_channels, num_classes, dropout=0.5):
        super().__init__()
        self.lin1 = Linear(hidden_channels, hidden_channels)
        self.lin2 = Linear(hidden_channels, num_classes)
        self.dropout = dropout

    def reset_parameters(self):
        self.lin1.reset_parameters()
        self.lin2.reset_parameters()

    def forward(self, x):
        x = F.relu(self.lin1(x))
        x = F.dropout(x, p=self.dropout, training=self.training)
        return self.lin2(x)


class _CLIPWrapper(nn.Module):

    def __init__(self, clip_model: nn.Module):
        super().__init__()
        self.clip_model = clip_model

    def forward(self, x):
        return self.clip_model.encode_text(x)


def conv3x3(in_planes: int, out_planes: int, stride: int=1, groups: int=1, dilation: int=1) ->nn.Conv2d:
    """3x3 convolution with padding."""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, groups=groups, bias=False, dilation=dilation)


class BasicBlock(nn.Module):
    expansion: int = 1
    __constants__ = ['downsample']

    def __init__(self, inplanes: int, planes: int, stride: int=1, downsample: Optional[nn.Module]=None, groups: int=1, base_width: int=64, dilation: int=1, norm_layer: Optional[Callable[..., nn.Module]]=None) ->None:
        super().__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if groups != 1 or base_width != 64:
            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
        if dilation > 1:
            raise NotImplementedError('Dilation > 1 not supported in BasicBlock')
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) ->Tensor:
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.downsample is not None:
            identity = self.downsample(x)
        out += identity
        out = self.relu(out)
        return out


def conv1x1(in_planes: int, out_planes: int, stride: int=1) ->nn.Conv2d:
    """1x1 convolution."""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class Bottleneck(nn.Module):
    expansion: int = 4
    __constants__ = ['downsample']

    def __init__(self, inplanes: int, planes: int, stride: int=1, downsample: Optional[nn.Module]=None, groups: int=1, base_width: int=64, dilation: int=1, norm_layer: Optional[Callable[..., nn.Module]]=None) ->None:
        super().__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.0)) * groups
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) ->Tensor:
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv3(out)
        out = self.bn3(out)
        if self.downsample is not None:
            identity = self.downsample(x)
        out += identity
        out = self.relu(out)
        return out


class ResNet(nn.Module):

    def __init__(self, block: Type[Union[BasicBlock, Bottleneck]], layers: List[int], zero_init_residual: bool=False, groups: int=1, widen: int=1, width_per_group: int=64, replace_stride_with_dilation: Optional[List[bool]]=None, norm_layer: Optional[Callable[..., nn.Module]]=None, first_conv3x3: bool=False, remove_first_maxpool: bool=False, in_chans: int=3) ->None:
        super().__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer
        self.inplanes = width_per_group * widen
        self.dilation = 1
        if replace_stride_with_dilation is None:
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError('replace_stride_with_dilation should be None or a 3-element tuple, got {}'.format(replace_stride_with_dilation))
        self.groups = groups
        self.base_width = width_per_group
        num_out_filters = width_per_group * widen
        if first_conv3x3:
            self.conv1 = nn.Conv2d(in_chans, num_out_filters, kernel_size=3, stride=1, padding=1, bias=False)
        else:
            self.conv1 = nn.Conv2d(in_chans, num_out_filters, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = norm_layer(num_out_filters)
        self.relu = nn.ReLU(inplace=True)
        if remove_first_maxpool:
            self.maxpool = nn.MaxPool2d(kernel_size=1, stride=1)
        else:
            self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, num_out_filters, layers[0])
        num_out_filters *= 2
        self.layer2 = self._make_layer(block, num_out_filters, layers[1], stride=2, dilate=replace_stride_with_dilation[0])
        num_out_filters *= 2
        self.layer3 = self._make_layer(block, num_out_filters, layers[2], stride=2, dilate=replace_stride_with_dilation[1])
        num_out_filters *= 2
        self.layer4 = self._make_layer(block, num_out_filters, layers[3], stride=2, dilate=replace_stride_with_dilation[2])
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)

    def _make_layer(self, block: Type[Union[BasicBlock, Bottleneck]], planes: int, blocks: int, stride: int=1, dilate: bool=False) ->nn.Sequential:
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(conv1x1(self.inplanes, planes * block.expansion, stride), norm_layer(planes * block.expansion))
        layers = [block(self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer)]
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes, groups=self.groups, base_width=self.base_width, dilation=self.dilation, norm_layer=norm_layer))
        return nn.Sequential(*layers)

    def forward(self, x: Tensor) ->Tensor:
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return x


class _VISSLBackboneWrapper(nn.Module):
    """VISSL backbones take additional arguments in ``forward`` that are not needed for our integration.

    This wrapper can be applied to a Flash backbone to ignore any additional arguments to ``forward``.
    """

    def __init__(self, backbone: nn.Module):
        super().__init__()
        self.backbone = backbone

    def forward(self, x, *args, **kwargs):
        x = self.backbone(x)
        x = x.unsqueeze(0)
        return x


class MockVISSLTask:
    """Mock task class from VISSL to support loss, configs, base_model, last batch etc."""

    def __init__(self, vissl_adapter, vissl_loss, task_config, vissl_model) ->None:
        self.vissl_adapter = vissl_adapter
        self.loss = vissl_loss
        self.config = task_config
        self.base_model = vissl_model
        self.model = self.base_model
        self.device = None
        self.iteration = 0
        self.max_iteration = 1
        self.last_batch = AttrDict({'sample': AttrDict({'input': None, 'data_momentum': None})})


class OnesModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(1, 2)
        self.register_buffer('zeros', torch.zeros(2))
        self.register_buffer('zero_one', torch.tensor([0.0, 1.0]))

    def forward(self, x):
        x = self.layer(x)
        return x * self.zeros + self.zero_one


class DummyClassifier(nn.Module):

    def __init__(self):
        super().__init__()
        self.backbone = nn.Sequential(nn.Flatten(), nn.Linear(28 * 28, 10))
        self.head = nn.LogSoftmax()

    def forward(self, x):
        return self.head(self.backbone(x))


import torch
from torch.nn import MSELoss, ReLU
from _paritybench_helpers import _mock_config, _mock_layer, _paritybench_base, _fails_compile


TESTCASES = [
    # (nn.Module, init_args, forward_args, jit_compiles)
    (BasicBlock,
     lambda: ([], {'inplanes': 4, 'planes': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (DefaultGraphHead,
     lambda: ([], {'hidden_channels': 4, 'num_classes': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (FuncModule,
     lambda: ([], {'func': _mock_layer()}),
     lambda: ([], {'input': torch.rand([4, 4])}),
     False),
    (_VISSLBackboneWrapper,
     lambda: ([], {'backbone': _mock_layer()}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
]

class Test_Lightning_AI_lightning_flash(_paritybench_base):
    def test_000(self):
        self._check(*TESTCASES[0])

    def test_001(self):
        self._check(*TESTCASES[1])

    def test_002(self):
        self._check(*TESTCASES[2])

    def test_003(self):
        self._check(*TESTCASES[3])

