code: def _ste_quantize(inputs, mean):
      """Calculates quantize(inputs - mean) + mean, sets straight-through grads."""
      half = tf.constant(0.5, dtype=tf.float32)
      outputs = inputs
      outputs -= mean
      # Rounding latents for the forward pass (straight-through).
      outputs = outputs + tf.stop_gradient(
          tf.math.floor(outputs + half) - outputs
      )
      outputs += mean
      return outputs
  
imports: import collections
  from compare_gan.architectures import abstract_arch
  from compare_gan.architectures import arch_ops
  import numpy
  import tensorflow.compat.v1
  import tensorflow_compression
  from hific.helpers import ModelMode
label_desc: Use image node_id 1 for calling a global function with example usage: _ste_quantize(inputs, mean) and returns: outputs