def blocksparse_matmul_grad()

in blocksparse/matmul.py [0:0]

• 34 lines of code
• 4 McCabe index (conditional complexity)

def blocksparse_matmul_grad(op, dy, temp):

    blocks    = op.get_attr("blocks")
    bsize     = op.get_attr("bsize")
    axis      = op.get_attr("axis")
    C         = op.get_attr("C")
    K         = op.get_attr("K")
    segments  = op.get_attr("segments_dx")
    shared    = op.get_attr("shared_dx")
    locks     = op.get_attr("locks_dx")
    gated_dw  = op.get_attr("gated_dw")
    gate_grad = op.get_attr("gate_grad")
    bench     = op.get_attr("bench")
    x         = op.inputs[0]
    w         = op.inputs[1]
    lut_dx    = op.inputs[3]
    lut_dw    = op.inputs[4]
    gate      = [op.inputs[5]] if len(op.inputs) > 5 else []
    name      = op.name.split('/')[-1]

    dx, _ = blocksparse_matmul_dx(
        dy, w, lut_dx, gate, gated_dw=gated_dw, gate_grad=gate_grad,
        blocks=blocks, bsize=bsize, axis=axis, C=K, K=C, # swap C,K
        segments=segments, locks=locks, shared=shared,
        bench=bench, name=name+"_bprop")

    dw = blocksparse_matmul_dw(
        [x], [dy], lut_dw, gate, gated_dw=gated_dw, gate_grad=gate_grad,
        blocks=blocks, bsize=bsize, axis=axis, C=C, K=K,
        bench=bench, name=name+"_updat")

    # print(dx.op.name, dx.op.device)
    # print(dw.op.name, dw.op.device)

    if len(gate) == 0:
        return (dx, dw, None, None, None)
    elif gate_grad:
        dw, dg = blocksparse_matmul_dg(dw, w, gate[0])
    else:
        dg = None

    return (dx, dw, None, None, None, dg)