def quantize()

in blocksparse/quantize.py [0:0]

• 40 lines of code
• 11 McCabe index (conditional complexity)

def quantize(x, qspec, b_qspec=None, name=None):

    if name is None:
        name = "quantize"

    if b_qspec is None:
        b_qspec = qspec

    if x.dtype.base_dtype == tf.bfloat16:
        for spec in (qspec, b_qspec):
            assert spec.fbits <= 7, "bfloat only supports up to 7 fractional bits"

    global log_init
    for spec in (qspec, b_qspec):
        if spec.logfile and spec.logfile not in log_init:
            with open(spec.logfile, 'w') as log:
                log.write("\t".join(quant_headers) + "\n")
            log_init.add(spec.logfile)

    e = [get_entropy()] if qspec.stoch == 2 else []

    reuse = tf.get_variable_scope().reuse

    with tf.device("/cpu:0"), tf.variable_scope("quantize"):
        exp_f = tf.get_variable(name + "_exp_f", dtype=tf.int64, initializer=np.int64(qspec.emax),   trainable=False)
        exp_b = tf.get_variable(name + "_exp_b", dtype=tf.int64, initializer=np.int64(b_qspec.emax), trainable=False)

    return quantize_op(x, exp_f, exp_b, e,
        ebits      = qspec.ebits,
        fbits      = qspec.fbits,
        stoch      = qspec.stoch,
        denorm     = qspec.denorm,
        freq       = (not reuse and qspec.freq),
        mode       = qspec.mode,
        bias_pad   = qspec.bias_pad,
        stdv_mul   = qspec.stdv_mul,
        logfile    = qspec.logfile,
        b_ebits    = b_qspec.ebits,
        b_fbits    = b_qspec.fbits,
        b_stoch    = b_qspec.stoch,
        b_denorm   = b_qspec.denorm,
        b_freq     = (not reuse and b_qspec.freq),
        b_mode     = b_qspec.mode,
        b_bias_pad = b_qspec.bias_pad,
        b_stdv_mul = b_qspec.stdv_mul,
        b_logfile  = b_qspec.logfile,
        name       = name,
    )