def block_reduced_full_dw()

in blocksparse/matmul.py [0:0]

• 31 lines of code
• 15 McCabe index (conditional complexity)

def block_reduced_full_dw(param_grad, scale=1.0, norm="max", group_size=8):

    # max(abs()) or l2_norm()
    norm  = 0 if norm.lower() == "max" else 1
    # host side scalar, if zero will cause compute for this op to be skipped.
    scale = scalar_constant(scale, dtype=tf.float32)

    assert group_size <= 8

    # backward walk param grad to find BlocksparseMatmulDW ops
    # this should only hit BlocksparseMatmulDWs, BlocksparseMatmulDGs, AddNs or FloatCasts
    ops = get_parents(param_grad, "BlocksparseMatmulDW")
    if len(ops) < 1:
        raise ValueError("BlocksparseMatmulDW op not found")

    # this sorting is dependent on the op names being correctly ordered.
    ops.sort(key=lambda op: op.name.split('/')[-1], reverse=True)

    # use the parent scope for the new ops
    scope = ops[-1].name.split('/')
    scope = '/'.join(scope[0:-1])

    # we're going to be using absolute names, so clear name_scope
    with tf.name_scope(None):
        dw_full = None
        offset  = 0
        while offset < len(ops):

            xs = [op.inputs[0] for op in ops[offset:offset+group_size] ]
            gs = [op.inputs[1] for op in ops[offset:offset+group_size] ]

            # Get the corresponding activation grad op for the last param grad op in the group
            bprop = None
            for consumer in gs[-1].consumers():
                if consumer.type == "BlocksparseMatmulDX":
                    bprop = consumer
                    break
            assert bprop is not None

            # get attributes of first op in group
            up    = ops[offset]
            bsize = up.get_attr("bsize")
            axis  = up.get_attr("axis")
            name  = "%s/block_reduced_full_dw_%03d" % (scope, offset)
            dw_full = [] if dw_full is None else [dw_full]

            dw_full, _, _ = blocksparse_reduced_dw(xs, gs, scale, dw_full, bsize=bsize, norm=norm, axis=axis, name=name)

            # force the dw op before any more time steps are processed
            bprop._add_control_input(dw_full.op)

            offset += group_size

    return dw_full