def _PendingCount()

in blocksparse/grads.py [0:0]

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

def _PendingCount(ys_ops, xs_ops):

    grad_dtypes = set((tf.float32, tf.float16, tf.bfloat16))

    # Ascend tree from the params and/or inputs (xs) to the losses (ys).
    # Create set of each unique node along the way.
    reached_ops = set()
    queue = collections.deque(xs_ops)
    while queue:
        op = queue.popleft()
        if op not in reached_ops:
            reached_ops.add(op)
            for output in op.outputs:
                if output.dtype.base_dtype in grad_dtypes:
                    queue.extend(output.consumers())
    # Get the subset of ys are reachable from xs.
    reachable_ys_ops = set(op for op in ys_ops if op in reached_ops)

    # Descend tree from ys along the reachable path.
    # Mark unique ops along the way (between_ops).
    # Handle gradient rerouting for recompute nodes.
    recompute_ops = list()
    between_ops   = set()
    queue         = collections.deque(reachable_ys_ops)
    while queue:
        op = queue.popleft()
        if op in reached_ops:
            between_ops.add(op)
            # don't add the inputs again.
            reached_ops.remove(op)
            # For recompute ops only traverse the second graph copy
            # We don't want the forward pass ops contributing to the pending_count.
            if op.type == "Recompute":
                recompute_ops.append(op)
                n_outs = len(op.outputs)
                for x in op.inputs[n_outs:n_outs*2]:
                    queue.append(x.op)
            else:
                for x in op.inputs:
                    queue.append(x.op)

    # Build a mapping from operation to the number of grad inputs to that op
    # ops not in this dict should no longer be traversed (excepting the initial ys ops with no dependancies).
    pending_count = dict()
    for op in between_ops:
        for x in op.inputs:
            if x.op in between_ops:
                pending_count[x.op] = pending_count.get(x.op, 0) + 1

    return pending_count, reachable_ys_ops, recompute_ops