def group_lstm_grads()

in blocksparse/lstm.py [0:0]

• 51 lines of code
• 14 McCabe index (conditional complexity)

def group_lstm_grads(grads, params, scope="grouped_lstm", group_size=None):

    grad = None
    grad_idx = None
    for i, (g, p) in enumerate(zip(grads, params)):
        if scope in p.name and "kernel" in p.name:
            grad = g
            grad_idx = i
            break
    assert grad is not None

    # backward walk param grad to find dw MatMul ops
    # walk should terminate with each MatMul op
    ops  = list()
    wave = set([grad.op])
    while wave:
        new_wave = set()
        for op in wave:
            for op in (t.op for t in op.inputs):
                # TN MatMul ops
                if op.type == "MatMul" and op.get_attr("transpose_a") and not op.get_attr("transpose_b"):
                    ops.append(op)
                else:
                    new_wave.add(op)
        wave = new_wave

    # sort op names descending and split out the lstms (if weights are shared)
    last_lstm = None
    lstms = list()
    ops.sort(key=lambda op: op.name, reverse=True)
    for op in ops:
        # gradients/grouped_lstm/lstm_2/step_00_grad/MatMul_1 => lstm_2
        lstm = op.name.split("/")[-3]
        if last_lstm != lstm:
            lstms.insert(0, list())
            last_lstm = lstm
        lstms[0].append(op)

    # we're going to be using absolute names, so clear name_scope
    with tf.name_scope(None):

        lstm_grads = list()
        for lstm_ops in lstms:

            # default dw op to one big matmul per lstm
            if group_size is None:
                group_size = len(lstm_ops)

            # use the lstm scope for the new ops
            # gradients/grouped_lstm/lstm_2/step_00_grad/MatMul_1 => gradients/grouped_lstm/lstm_2
            scope = lstm_ops[-1].name.split('/')
            scope = '/'.join(scope[0:-2])

            offset = 0
            while offset < len(lstm_ops):

                xs = tf.concat([op.inputs[0] for op in lstm_ops[offset:offset+group_size] ], axis=0)
                gs = tf.concat([op.inputs[1] for op in lstm_ops[offset:offset+group_size] ], axis=0)

                mmop = tf.matmul(xs, gs, transpose_a=True, transpose_b=False, name="%s/dw_%04d" % (scope, offset))
                grad = mmop if offset == 0 else ew.add(grad, mmop, name="%s/add_%04d" % (scope, offset))

                offset += group_size

            lstm_grads.append(grad)

        if len(lstms) > 1:
            from blocksparse.ewops import add_n
            # gradients/grouped_lstm/lstm_2/step_00_grad/MatMul_1 => gradients/grouped_lstm
            scope = lstms[0][-1].name.split('/')
            scope = '/'.join(scope[0:-3])
            grads[grad_idx] = tf.add_n(lstm_grads, name="%s/add_n" % scope)
        else:
            grads[grad_idx] = lstm_grads[0]