class_docstr: Straight-Through Estimation (STE) function.
  
  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:
      grad_fn: function to restrain the gradient received. If None, no mapping will performed.
  
  Example:
      Let the gradients of ``torch.sign`` estimated from STE.
      >>> input = torch.randn(4, requires_grad = True)
      >>> output = torch.sign(input)
      >>> loss = output.mean()
      >>> loss.backward()
      >>> input.grad
      tensor([0., 0., 0., 0.])
  
      >>> with torch.no_grad():
      ...     output = torch.sign(input)
      >>> out_est = STEFunction.apply(input, output)
      >>> loss = out_est.mean()
      >>> loss.backward()
      >>> input.grad
      tensor([0.2500, 0.2500, 0.2500, 0.2500])
class_docstr_tok: ["Straight-Through","Estimation","(","STE",")","function",".","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",":","grad_fn",":","function","to","restrain","the","gradient","received",".","If","None",",","no","mapping","will","performed",".","Example",":","Let","the","gradients","of","``","torch.sign","``","estimated","from","STE",".",">",">",">","input","=","torch.randn","(","4",",","requires_grad","=","True",")",">",">",">","output","=","torch.sign","(","input",")",">",">",">","loss","=","output.mean","(",")",">",">",">","loss.backward","(",")",">",">",">","input.grad","tensor","(","[","0.",",","0.",",","0.",",","0",".","]",")",">",">",">","with","torch.no_grad","(",")",":","...","output","=","torch.sign","(","input",")",">",">",">","out_est","=","STEFunction.apply","(","input",",","output",")",">",">",">","loss","=","out_est.mean","(",")",">",">",">","loss.backward","(",")",">",">",">","input.grad","tensor","(","[","0.2500",",","0.2500",",","0.2500",",","0.2500","]",")"]
code: def backward(ctx: Any, grad_output: Tensor) -> Tuple[Tensor, Tensor, None]:  # type: ignore[override]
      if ctx.grad_fn is None:
          return (
              grad_output.sum_to_size(ctx.in_shape),
              grad_output.sum_to_size(ctx.out_shape),
              None,
          )
      return (
          ctx.grad_fn(grad_output.sum_to_size(ctx.in_shape)),
          ctx.grad_fn(grad_output.sum_to_size(ctx.out_shape)),
          None,
      )
  
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 2 for calling the STEFunction obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: ctx,  ctx,  None, grad_output,  grad_output,  None