class_docstr: Straight-Through Estimation (STE) module.
  
  STE wraps the ``STEFunction`` to aid the back propagation of non-differentiable modules.
  It may also use to avoid gradient computation for differentiable operations. By default,
  STE bridges the gradients between the input tensor and output tensor as if the function
  was an identity function. Meanwhile, advanced gradient functions are also supported. e.g.
  the output gradients can be mapped into [-1, 1] with ``F.hardtanh`` function.
  
  Args:
      target_fn: the target function to wrap with.
      grad_fn: function to restrain the gradient received. If None, no mapping will performed.
  
  Example:
      ``RandomPosterize`` is a non-differentiable operation. Let STE estimate the gradients.
      >>> import kornia.augmentation as K
      >>> import torch.nn.functional as F
      >>> input = torch.randn(1, 1, 4, 4, requires_grad = True)
      >>> estimator = StraightThroughEstimator(K.RandomPosterize(3, p=1.), grad_fn=F.hardtanh)
      >>> out = estimator(input)
      >>> out.mean().backward()
      >>> input.grad
      tensor([[[[0.0625, 0.0625, 0.0625, 0.0625],
                [0.0625, 0.0625, 0.0625, 0.0625],
                [0.0625, 0.0625, 0.0625, 0.0625],
                [0.0625, 0.0625, 0.0625, 0.0625]]]])
  
      This can be used to chain up the gradients within a ``Sequential`` block.
      >>> import kornia.augmentation as K
      >>> input = torch.randn(1, 1, 4, 4, requires_grad = True)
      >>> aug = K.ImageSequential(
      ...     K.RandomAffine((77, 77)),
      ...     StraightThroughEstimator(K.RandomPosterize(3, p=1.), grad_fn=None),
      ...     K.RandomRotation((15, 15)),
      ... )
      >>> aug(input).mean().backward()
      >>> input.grad
      tensor([[[[0.0422, 0.0626, 0.0566, 0.0422],
                [0.0566, 0.0626, 0.0626, 0.0626],
                [0.0626, 0.0626, 0.0626, 0.0566],
                [0.0422, 0.0566, 0.0626, 0.0422]]]])
class_docstr_tok: ["Straight-Through","Estimation","(","STE",")","module",".","STE","wraps","the","``","STEFunction","``","to","aid","the","back","propagation","of","non-differentiable","modules",".","It","may","also","use","to","avoid","gradient","computation","for","differentiable","operations",".","By","default",",","STE","bridges","the","gradients","between","the","input","tensor","and","output","tensor","as","if","the","function","was","an","identity","function",".","Meanwhile",",","advanced","gradient","functions","are","also","supported",".","e.g",".","the","output","gradients","can","be","mapped","into","[","-1",",","1","]","with","``","F.hardtanh","``","function",".","Args",":","target_fn",":","the","target","function","to","wrap","with",".","grad_fn",":","function","to","restrain","the","gradient","received",".","If","None",",","no","mapping","will","performed",".","Example",":","``","RandomPosterize","``","is","a","non-differentiable","operation",".","Let","STE","estimate","the","gradients",".",">",">",">","import","kornia.augmentation","as","K",">",">",">","import","torch.nn.functional","as","F",">",">",">","input","=","torch.randn","(","1",",","1",",","4",",","4",",","requires_grad","=","True",")",">",">",">","estimator","=","StraightThroughEstimator","(","K.RandomPosterize","(","3",",","p=1",".",")",",","grad_fn=F.hardtanh",")",">",">",">","out","=","estimator","(","input",")",">",">",">","out.mean","(",")",".backward","(",")",">",">",">","input.grad","tensor","(","[","[","[","[","0.0625",",","0.0625",",","0.0625",",","0.0625","]",",","[","0.0625",",","0.0625",",","0.0625",",","0.0625","]",",","[","0.0625",",","0.0625",",","0.0625",",","0.0625","]",",","[","0.0625",",","0.0625",",","0.0625",",","0.0625","]","]","]","]",")","This","can","be","used","to","chain","up","the","gradients","within","a","``","Sequential","``","block",".",">",">",">","import","kornia.augmentation","as","K",">",">",">","input","=","torch.randn","(","1",",","1",",","4",",","4",",","requires_grad","=","True",")",">",">",">","aug","=","K.ImageSequential","(","...","K.RandomAffine","(","(","77",",","77",")",")",",","...","StraightThroughEstimator","(","K.RandomPosterize","(","3",",","p=1",".",")",",","grad_fn=None",")",",","...","K.RandomRotation","(","(","15",",","15",")",")",",","...",")",">",">",">","aug","(","input",")",".mean","(",")",".backward","(",")",">",">",">","input.grad","tensor","(","[","[","[","[","0.0422",",","0.0626",",","0.0566",",","0.0422","]",",","[","0.0566",",","0.0626",",","0.0626",",","0.0626","]",",","[","0.0626",",","0.0626",",","0.0626",",","0.0566","]",",","[","0.0422",",","0.0566",",","0.0626",",","0.0422","]","]","]","]",")"]
code: def forward(self, input: Tensor) -> Tensor:
      out = self.target_fn(input)
      if not isinstance(out, Tensor):
          raise NotImplementedError(
              "Only Tensor is supported at the moment. Feel free to contribute to https://github.com/kornia/kornia."
          )
      output = STEFunction.apply(input, out, self.grad_fn)
      return output
  
file_path: kornia/kornia/grad_estimator/ste.py
imports: from typing import Any, Callable, Optional, Tuple
  from torch import Tensor, nn
  from torch.autograd import Function
label_desc: Use image node_id 3 for calling the StraightThroughEstimator obj's underlying member method code with example usage: obj.forward(input) and returns: output