class_docstr: A universal tensor quantization function
  
  Take an input tensor, output an quantized tensor. The granularity of scale can be interpreted from the
  shape of amax.
  output_dtype indicates whether the quantized value will be stored in integer or float. The reason we want to store
  it in float is the pytorch function takes the quantized value may not accept integer input, e.g. Conv2D.
  
  It uses 2^num_bits -1 values instead of 2^num_bits. e.g., for num_bits=8, it uses [-127, 127] instead of [-128, 127]
class_docstr_tok: ["A","universal","tensor","quantization","function","Take","an","input","tensor",",","output","an","quantized","tensor",".","The","granularity","of","scale","can","be","interpreted","from","the","shape","of","amax",".","output_dtype","indicates","whether","the","quantized","value","will","be","stored","in","integer","or","float",".","The","reason","we","want","to","store","it","in","float","is","the","pytorch","function","takes","the","quantized","value","may","not","accept","integer","input",",","e.g",".","Conv2D",".","It","uses","2^num_bits","-1","values","instead","of","2^num_bits",".","e.g.",",","for","num_bits=8",",","it","uses","[","-127",",","127","]","instead","of","[","-128",",","127","]"]
docstr: Implements straight through estimation with clipping. For -amax <= input <= amax
  the gradient passes straight through, otherwise the gradient is zero.
  
  Args:
      ctx: A Context object with saved tensors from forward.
      grad_outputs: A tensor of gradient of outputs.
      grad_scale: A tensor of gradient of scale.
  
  Returns:
      grad_inputs: A tensor of gradient.
docstr_tok: ["Implements","straight","through","estimation","with","clipping",".","For","-amax","<","=","input","<","=","amax","the","gradient","passes","straight","through",",","otherwise","the","gradient","is","zero",".","Args",":","ctx",":","A","Context","object","with","saved","tensors","from","forward",".","grad_outputs",":","A","tensor","of","gradient","of","outputs",".","grad_scale",":","A","tensor","of","gradient","of","scale",".","Returns",":","grad_inputs",":","A","tensor","of","gradient","."]
code: def backward(ctx, grad_outputs, grad_scale):
      """
      Implements straight through estimation with clipping. For -amax <= input <= amax
      the gradient passes straight through, otherwise the gradient is zero.
  
      Args:
          ctx: A Context object with saved tensors from forward.
          grad_outputs: A tensor of gradient of outputs.
          grad_scale: A tensor of gradient of scale.
  
      Returns:
          grad_inputs: A tensor of gradient.
      """
      inputs, amax = ctx.saved_tensors
      zero = grad_outputs.new_zeros(
          1
      )  # create a zero tensor with the same type and device
      grad_inputs = torch.where(
          inputs.abs() <= amax, grad_outputs, zero
      )
      return grad_inputs, None, None, None, None
  
imports: import numpy
  import yaml
  import logging
  import torch
  from torch.autograd import Function
label_desc: Use image node_id 2 for calling the TensorQuantFunction obj's underlying member method code with example usage: obj.backward(ctx, grad_outputs, grad_scale) and returns: grad_inputs,  None,  None,  None,  None