sequenceDiagram
    participant "AttributionCommon"
    participant "init_attr_cache_in (callback)"
    participant "TensorDict"
    participant "additional_init_tensors"
    "AttributionCommon"->>"TensorDict": Get cache_in
    alt init_attr_cache_in is provided
        "AttributionCommon"->>"init_attr_cache_in (callback)": Call with cache_in, additional_init_tensors
        "init_attr_cache_in (callback)"-->>"AttributionCommon": Return new cache_in
        "AttributionCommon"->>"TensorDict": Validate type of cache_in
    end
    "AttributionCommon"->>"TensorDict": Continue attribution process with cache_in

classDiagram
    class AttributionCommon {
        - target_modules: Optional[List[str]]
        - init_attr_targets: Optional[Callable]
        - init_attr_inputs: Optional[Callable]
        - init_attr_cache_in: Optional[Callable]
        - init_attr_grads: Optional[Callable]
        - additional_init_keys: Optional[List[UnraveledKey]]
        - output_grad_callbacks: Optional[Dict[str, Callable]]
        - cache_callback: Optional[Callable]
        ...
    }
    class GradCAM {
        - target_modules: Optional[List[str]]
        - init_attr_targets: Optional[Callable]
        - init_attr_inputs: Optional[Callable]
        - init_attr_cache_in: Optional[Callable]
        - init_attr_grads: Optional[Callable]
        - additional_init_keys: Optional[List[UnraveledKey]]
        - output_grad_callbacks: Optional[Dict[str, Callable]]
        ...
    }
    class GuidedBackpropagation {
        - target_modules: Optional[List[str]]
        - init_attr_targets: Optional[Callable]
        - init_attr_inputs: Optional[Callable]
        - init_attr_cache_in: Optional[Callable]
        - init_attr_grads: Optional[Callable]
        - additional_init_keys: Optional[List[UnraveledKey]]
        - output_grad_callbacks: Optional[Dict[str, Callable]]
        ...
    }
    class IntegratedGradients {
        - target_modules: Optional[List[str]]
        - init_attr_targets: Optional[Callable]
        - init_attr_inputs: Optional[Callable]
        - init_attr_cache_in: Optional[Callable]
        - init_attr_grads: Optional[Callable]
        - additional_init_keys: Optional[List[UnraveledKey]]
        - output_grad_callbacks: Optional[Dict[str, Callable]]
        ...
    }
    class LRP {
        - target_modules: Optional[List[str]]
        - init_attr_targets: Optional[Callable]
        - init_attr_inputs: Optional[Callable]
        - init_attr_cache_in: Optional[Callable]
        - init_attr_grads: Optional[Callable]
        - additional_init_keys: Optional[List[UnraveledKey]]
        - output_grad_callbacks: Optional[Dict[str, Callable]]
        ...
    }
    class Saliency {
        - target_modules: Optional[List[str]]
        - init_attr_targets: Optional[Callable]
        - init_attr_inputs: Optional[Callable]
        - init_attr_cache_in: Optional[Callable]
        - init_attr_grads: Optional[Callable]
        - additional_init_keys: Optional[List[UnraveledKey]]
        - output_grad_callbacks: Optional[Dict[str, Callable]]
        ...
    }
    GradCAM --|> AttributionCommon
    GuidedBackpropagation --|> AttributionCommon
    IntegratedGradients --|> AttributionCommon
    LRP --|> AttributionCommon
    Saliency --|> AttributionCommon