def group_dg_grads()

in blocksparse/matmul.py [0:0]

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

def group_dg_grads(bsmm_dw_op, dw, scope):

    # splice the dg + addn ops out of the graph and replace with a single dg op
    # that takes in the final accumulated dw value
    dg_op  = bsmm_dw_op.outputs[0].consumers()[0]
    assert dg_op.type == "BlocksparseMatmulDG"
    dw, dg = blocksparse_matmul_dg(dw, *dg_op.inputs[1:], name=f"{scope}/BlocksparseMatmulDG")

    # splice old add_n op out of graph
    addn_op  = dg_op.outputs[1].consumers()[0]
    addn_ops = list()
    addn_ops.append(addn_op)
    if addn_op.type[0:3] != "Add":
        raise ValueError(f"bad type: {addn_ops[0].type} Cause: this segment does not share a broadcasted gate.")
    elif addn_op.type == "AddN8":
        while True:
            addn_op = addn_op.outputs[0].consumers()[0]
            if addn_op.type == "AddN8":
                addn_ops.append(addn_op)
            else:
                break

    # print(addn_op.name)
    # for i in addn_op.inputs:
    #     print(i.name)
    # print()
    addn = addn_ops[-1].outputs[0]
    dg_consumers = addn.consumers()
    #for op in dg_consumers:

    assert len(dg_consumers) > 0, "raw dg grad not supported"
    #print(addn.name)
    for dg_consumer in dg_consumers:
        found = False
        #print(dg_consumer.name)
        for i, t in enumerate(dg_consumer.inputs):
            #print(i, t.name)
            if t is addn:
                #print(f"splicing dg into: {dg_consumer.name} at {i}")
                dg_consumer._update_input(i, dg)
                found = True
                break
        if not found:
            print(f"splice failed for {dg_consumer.name}")
    return dw