class_docstr: Base class of differentiable augmentation operations.
  
  Args:
      operation: Kornia augmentation module.
      initial_magnitude: targeted magnitude parameter name and its initial magnitude value.
          The magnitude parameter name shall align with the attribute inside the random_generator
          in each augmentation. If None, the augmentation will be randomly applied according to
          the augmentation sampling range.
      temperature: temperature for RelaxedBernoulli distribution used during training.
      is_batch_operation: determine if to obtain the probability from `p` or `p_batch`.
          Set to True for most non-shape-persistent operations (e.g. cropping).
class_docstr_tok: ["Base","class","of","differentiable","augmentation","operations",".","Args",":","operation",":","Kornia","augmentation","module",".","initial_magnitude",":","targeted","magnitude","parameter","name","and","its","initial","magnitude","value",".","The","magnitude","parameter","name","shall","align","with","the","attribute","inside","the","random_generator","in","each","augmentation",".","If","None",",","the","augmentation","will","be","randomly","applied","according","to","the","augmentation","sampling","range",".","temperature",":","temperature","for","RelaxedBernoulli","distribution","used","during","training",".","is_batch_operation",":","determine","if","to","obtain","the","probability","from","`","p","`","or","`","p_batch","`",".","Set","to","True","for","most","non-shape-persistent","operations","(","e.g",".","cropping",")","."]
code: def _init_magnitude_fn(
      self, magnitude_fn: Optional[Callable[[Tensor], Tensor]]
  ) -> Callable[[Tensor], Tensor]:
      def _identity(x: Tensor) -> Tensor:
          return x
  
      def _random_flip(
          fn: Callable[[Tensor], Tensor]
      ) -> Callable[[Tensor], Tensor]:
          def f(x: Tensor) -> Tensor:
              flip = torch.rand((x.shape[0],), device=x.device) > 0.5
              return fn(x) * flip
  
          return f
  
      if magnitude_fn is None:
          magnitude_fn = _identity
  
      if self.symmetric_megnitude:
          return _random_flip(magnitude_fn)
  
      return magnitude_fn
  
imports: from typing import Callable, Dict, List, Optional, Tuple, Type, TypeVar
  import torch
  from torch import nn
  from torch.autograd import Function
  from torch.distributions import Bernoulli, RelaxedBernoulli
  from kornia.augmentation.base import _AugmentationBase
  from kornia.core import Module, Tensor
label_desc: Use image node_id 2 for calling the OperationBase obj's underlying member method code with example usage: obj._init_magnitude_fn(magnitude_fn) and returns: magnitude_fn, x, f, _random_flip, unknown