def apply()

in blocksparse/optimize.py [0:0]

• 23 lines of code
• 10 McCabe index (conditional complexity)

    def apply(self, params, qspec=None):

        with tf.device("/gpu:0"), tf.control_dependencies(None):
            for param in params:
                if self.fp16 == 2 or (self.fp16 and is_param_casted(param)):
                    # only use fp16 for params that are explicitly cast to fp16 before use
                    init  = float_cast(param.initialized_value(), dtype=tf.float16)
                    dtype = tf.float16
                else:
                    init  = param.initialized_value()
                    dtype = tf.float32

                with tf.variable_scope(None, param.op.name + "/" + self.name):
                    # use the Identity read op output as the key
                    # this lets us lookup ema vars by Cast op outputs
                    self.averages[param.value()] = tf.get_variable("ema", dtype=dtype, initializer=init, trainable=False)
                ops.add_to_collection(ops.GraphKeys.MOVING_AVERAGE_VARIABLES, param)

        ema_ops = []
        for param in params:

            ema   = self.averages[param.value()]
            gate  = getattr(param, "gate", None)
            gate  = [gate] if self.gated and gate is not None else []

            op = ema_op(ema, param, gate, decay=self.decay)

            if qspec is not None:
                ema_ops.append(ema.assign(quantize(op, qspec, name="ema_" + param.op.name)))
            else:
                ema_ops.append(op)

        return tf.group(*ema_ops)