code: def _threshold_mask(
      self,
      mask: Optional[Tensor],
      threshold_config: ThresholdConfig,
  ) -> Optional[Tensor]:
      if mask is None:
          return None
  
      if threshold_config.type == ThresholdType.hard:
          return (mask > threshold_config.value).float()
  
      if threshold_config.type in [
          ThresholdType.topk,
          ThresholdType.topk_hard,
      ]:
          if threshold_config.value >= mask.numel():
              if threshold_config.type == ThresholdType.topk:
                  return mask
              else:
                  return torch.ones_like(mask)
  
          value, index = torch.topk(
              mask.flatten(),
              k=threshold_config.value,
          )
  
          out = torch.zeros_like(mask.flatten())
          if threshold_config.type == ThresholdType.topk:
              out[index] = value
          else:
              out[index] = 1.0
          return out.view(mask.size())
  
      assert False
  
imports: import copy
  from typing import Dict, List, Optional, Union
  import torch
  from torch import Tensor
  from torch_geometric.data.data import Data, warn_or_raise
  from torch_geometric.data.hetero_data import HeteroData
  from torch_geometric.explain.config import ThresholdConfig, ThresholdType
  from torch_geometric.typing import EdgeType, NodeType
  from torch_geometric.visualization import visualize_graph
label_desc: Use image node_id 3 for calling the ExplanationMixin obj's underlying member method code with example usage: obj._threshold_mask(mask, threshold_config) and returns: None, unknown, out, mask, torch