//  [python]
def quantize(inputs, medians=None, method="SGAo-Q", **params):
    if method == "Identity":
        # (inputs - medians) + medians = inputs
        return inputs

    if method == "AUN-Q":
        half = tf.constant(0.5)
        noise = tf.random.uniform(tf.shape(inputs), -half, half)
        # (inputs - medians) + noise + medians = inputs + noise
        outputs = tf.math.add_n([inputs, noise])
        return outputs

    if medians is not None:
        inputs_ = inputs - medians
    else:
        inputs_ = inputs

    if method == "STE-Q":
        outputs = tf.round(inputs_)

        if medians is not None:
            outputs += medians
        outputs = tf.stop_gradient(outputs - inputs) + inputs

    elif method in {"St-Q", "SRA-Q"}:
        diff = inputs_ - tf.floor(inputs_)

        if method == "St-Q":
            probability = diff
        else:
            tau = params["tau"]
            likelihood_up = tf.exp(-tf.atanh(diff) / tau)
            likelihood_down = tf.exp(-tf.atanh(1 - diff) / tau)
            probability = likelihood_down / (likelihood_up + likelihood_down)
        delta = tf.to_float(
            (probability >= tf.random.uniform(tf.shape(probability))),
        )
        outputs = tf.floor(inputs_) + delta
        if medians is not None:
            outputs += medians
        outputs = tf.stop_gradient(outputs - inputs) + inputs

    elif method == "U-Q":
        half = tf.constant(0.5)
        # random value, shape: (N, 1, 1, 1)
        noise = tf.random.uniform(tf.shape(inputs), -half, half)[
            :, 0:1, 0:1, 0:1
        ]
        outputs = tf.round(inputs_ + noise) - noise
        if medians is not None:
            outputs += medians
        outputs = tf.stop_gradient(outputs - inputs) + inputs
        # tf.round(inputs - medians + noise) - noise + medians
        # = tf.round(inputs + nmm) - nmm, nmm = noise minus medians
    elif method == "SGA-Q":
        tau = params["tau2"]
        # use Gumbel Softmax implemented in tfp.distributions.RelaxedOneHotCategorical
        # The code is modified from https://github.com/mandt-lab/improving-inference-for-neural-image-compression/blob/main/sga.py#L110-L121
        # copyright @yiboyang
        epsilon = 1e-5
        y_floor = tf.floor(inputs_)
        y_ceil = tf.ceil(inputs_)
        y_bds = tf.stack([y_floor, y_ceil], axis=-1)
        ry_logits = tf.stack(
            [
                -tf.math.atanh(
                    tf.clip_by_value(inputs_ - y_floor, -1 + epsilon, 1 - epsilon)
                )
                / tau,
                -tf.math.atanh(
                    tf.clip_by_value(y_ceil - inputs_, -1 + epsilon, 1 - epsilon)
                )
                / tau,
            ],
            axis=-1,
        )
        # last dim are logits for DOWN or UP
        ry_dist = tfp.distributions.RelaxedOneHotCategorical(tau, logits=ry_logits)
        ry_sample = ry_dist.sample()
        outputs = tf.reduce_sum(y_bds * ry_sample, axis=-1)
        if medians is not None:
            outputs += medians

    elif method == "sigmoid":
        T = params["T"]
        diff = inputs_ - tf.floor(inputs_)
        temp = diff * T
        # following the original implementation
        # https://github.com/aliyun/alibabacloud-quantization-networks/blob/master/anybit.py#L26
        temp = tf.clip_by_value(temp, -10.0, 10.0)
        outputs = tf.sigmoid(temp) + tf.floor(inputs_)
        if medians is not None:
            outputs += medians

    elif method == "DS-Q":
        k = params["k"]
        y_floor = tf.floor(inputs_)
        # 0 <= diff <= 1
        diff = inputs_ - y_floor
        # -1 <= phi <= 1
        phi = tf.tanh((diff - 0.5) * k) / tf.tanh(0.5 * k)
        # y_floor <= y_phi <= y_floor + 1
        y_phi = (1 + phi) / 2 + y_floor
        outputs = tf.round(inputs_)
        outputs = tf.stop_gradient(outputs - y_phi) + y_phi
        if medians is not None:
            outputs += medians
    else:
        raise NotImplementedError
    return outputs