## @package control_ops_grad
# Module caffe2.python.control_ops_grad





from caffe2.proto import caffe2_pb2


def gen_do_gradient(op, g_output):
    """
    Generates gradient Do operator, given forward Do op and a list
    of gradient blobs corresponding to forward op's outputs
    Returns a gradient op and a list of blobs corresponding to input gradients
    """
    from caffe2.python.core import BlobReference
    subnet, outer_to_inner_map, inner_to_outer_map, workspace_blob_name = \
        _do_op_sanity_check_and_process(op)

    assert len(g_output) == len(op.output), \
        "Different number of gradient blobs and Do op outputs"

    grad_ops, deduped_g_output = dedupe_g_output(op, g_output)
    g_output = deduped_g_output

    # From the outer net point of view:
    #  Do is an operator that has some number of inputs and outputs;
    #  we have to generate a gradient operator that writes into
    #  corresponding input gradient blobs and has access to inputs, outputs
    #  and gradient output blobs
    # From the inner net point of view:
    #  Do is an operator with a subnet and blob bindings,
    #  we need to forward Do's output blob gradients into inner workspace,
    #  use them to run backward pass generation and forward Do's input blob
    #  gradients back into outer workspace

    op_output = [str(o) for o in op.output]
    op_output = op_output[:-1]  # remove workspace pointer blob
    op_input = [str(i) for i in op.input]
    op_input = op_input[:-1]  # remove workspace pointer blob

    ordered_inner_output_blob_names = [outer_to_inner_map[o] for o in op_output]

    backward_pass_initial_grad_map = {}
    initial_grad_map = {}
    for inner_output_name, outer_grad_output_name in \
            zip(ordered_inner_output_blob_names, g_output):
        # link inner_output_name to corresponding inner_grad_output_name for
        # backward pass generation;
        if outer_grad_output_name:
            inner_grad_output_name = inner_output_name + "/_DO_OPERATOR_INNER_GRAD_"
            backward_pass_initial_grad_map[BlobReference(inner_output_name)] = \
                BlobReference(inner_grad_output_name)
            initial_grad_map[inner_grad_output_name] = str(outer_grad_output_name)
    assert len(initial_grad_map) > 0, "Empty initial gradient map for Do op"

    inner_grad_ops, inner_grad_names_map = _gen_subgradient_pass(
        subnet, backward_pass_initial_grad_map)

    if len(inner_grad_ops) == 0:
        return [], []

    grad_copy_ops = []
    g_input = []
    new_op_outputs = []
    new_blob_bindings = {}
    for outer_input_name in op_input:
        inner_input_name = outer_to_inner_map[outer_input_name]
        if inner_input_name in inner_grad_names_map:
            inner_grad_input_name = inner_grad_names_map[inner_input_name]
            outer_grad_input_name = outer_input_name + "_grad"

            # It is possible that inner_grad_input_name will need to be
            # linked to another outer blob. For example:
            #
            #    // y - param initialized in init_net
            #    x = ...
            #    z = ...
            #    with ops.IfNet(...):
            #        ops.Add([z, x], y) # inner Do block
            #    loss = f(..., y, ...)
            #
            # In this case x, y and z are external for the inner Do block,
            # the inputs of the Do block are z and x and the output is y.
            # When computing the gradient of input x given the gradient
            # of output y it's easy to see that they are equal.
            # During the generation of gradient Do operator, we link
            # external gradient y (y_grad) to the internal name
            # (y/_DO_OPERATOR_INNER_GRAD_) and generate the backward pass
            # for the internal Do net. As a result we get gradient operators
            # for the gradient Do and gradient map that maps internal Do
            # blobs to their computed gradients.
            # In this example, gradient map may have blob x linked to
            # gradient blob y/_DO_OPERATOR_INNER_GRAD_.
            # We should export gradient for x outside of Do, so
            # we add a blob mapping from inner gradient blob
            # (y/_DO_OPERATOR_INNER_GRAD_) to a new outer name (x_grad).
            #
            # (Note: since we use transparent blob mapping between outer and
            # inner (Do's) workspace, these operations do not involve copying
            # but are merely using blobs in outer workspace in the Do's operator
            # workspace under (possibly) different names)
            #
            # At the same time, we need to add a blob mapping from inner name
            # y/_DO_OPERATOR_INNER_GRAD_ to the outer blob y_grad
            # Hence in this case, we cannot use existing blob mapping scheme
            # that requires a bijection between subset of inner blob names and
            # a set of all (Do's input and output) outer blob names

            # TODO(iliacher): Remove unnecessary blob copying

            new_inner_grad_input_name = \
                inner_input_name + "/_DO_OPERATOR_INNER_GRAD_COPY_"
            grad_copy_ops.append(_prepare_blob_copy_op(
                inner_grad_input_name, new_inner_grad_input_name))

            new_blob_bindings[new_inner_grad_input_name] = outer_grad_input_name
            new_op_outputs.append(outer_grad_input_name)
            g_input.append(outer_grad_input_name)
        else:
            g_input.append(None)

    new_op_inputs = []
    overwritten_names = set()
    saved_local_blob_names = set()
    for grad_op in inner_grad_ops:
        grad_op_input = [str(i) for i in grad_op.input]
        grad_op_output = [str(o) for o in grad_op.output]
        for grad_op_input_name in grad_op_input:
            if grad_op_input_name in overwritten_names:
                continue
            # check if this is an external blob
            outer_name = inner_to_outer_map.get(grad_op_input_name, None)
            if not outer_name:
                # check if this is an external gradient blob
                outer_name = initial_grad_map.get(grad_op_input_name, None)
            if outer_name:
                outer_name = str(outer_name)
                if outer_name not in new_op_inputs:
                    new_op_inputs.append(outer_name)

                new_blob_bindings[grad_op_input_name] = outer_name
            else:
                # this is a local blob, we'll get it's value from
                # a saved forward op workspace
                saved_local_blob_names.add(grad_op_input_name)
        overwritten_names.update(grad_op_output)

    # add inner gradient copy ops
    inner_grad_ops += grad_copy_ops

    gradient_do_def = _prepare_gradient_do_op(
        fwd_op=op,
        fwd_net=subnet,
        grad_ops=inner_grad_ops,
        inputs=new_op_inputs,
        outputs=new_op_outputs,
        blob_bindings=new_blob_bindings,
        saved_fwd_blobs=saved_local_blob_names,
        workspace_blob_name=workspace_blob_name)
    grad_ops.append(gradient_do_def)

    _do_op_sanity_check_and_process(gradient_do_def)

    return grad_ops, g_input


def dedupe_g_output(op, g_output):
    # When generation a gradient op it's possible to receive the same gradient
    # blob corresponding to different forward op output blobs, Do operator
    # requires a bijection between inner and outer names, make sure we do
    # deduplication
    grad_ops = []
    deduped_g_output = []
    init_grad_map = {}
    for output_name, grad_name in zip(op.output, g_output):
        if not grad_name:
            deduped_g_output.append(grad_name)
            continue

        if output_name in init_grad_map:
            deduped_g_output.append(init_grad_map[output_name])
        else:
            if grad_name not in init_grad_map.values():
                init_grad_map[output_name] = grad_name
                deduped_g_output.append(grad_name)
            else:
                deduped_grad_name = output_name + "_" + grad_name + "_DEDUP"
                assert deduped_grad_name not in init_grad_map.values()
                grad_copy_op = caffe2_pb2.OperatorDef()
                grad_copy_op.type = "Copy"
                grad_copy_op.input.extend([grad_name])
                grad_copy_op.output.extend([deduped_grad_name])
                grad_ops.append(grad_copy_op)
                deduped_g_output.append(deduped_grad_name)
                init_grad_map[output_name] = deduped_grad_name
    return grad_ops, deduped_g_output


def gen_while_gradient(op, g_output):
    """
    Generates gradient While operator
    """
    from caffe2.python.core import BlobReference
    assert op.type == "While", "Expected While op"
    assert len(op.input) > 0, "Expected at least one input in While op"

    assert len(op.output) == len(g_output), \
        "Different number of gradient blobs and While op outputs"

    grad_ops, deduped_g_output = dedupe_g_output(op, g_output)
    g_output = deduped_g_output

    init_grad_map = {}
    op_output = [str(o) for o in op.output]
    for output_name, grad_output_name in zip(op_output, g_output):
        if grad_output_name:
            init_grad_map[BlobReference(output_name)] = \
                BlobReference(grad_output_name)
    assert len(init_grad_map) > 0, "Empty initial gradient map for While op"

    loop_net = _get_net_argument(op, "loop_net")
    assert loop_net, "Expected loop subnet in While op"
    assert len(loop_net.op) == 1 and loop_net.op[0].type == "Do", \
        "Gradient While op requires single Do op as a loop body"
    do_op = loop_net.op[0]
    do_args = _get_do_arguments(do_op)
    assert "reuse_workspace" not in do_args or not do_args["reuse_workspace"], \
        "Gradient While op requires Do loop body op without reuse_workspace set"

    assert len(do_op.output) > 0, "Expected Do op with at least one output"
    workspace_blob = do_op.output[-1]

    loop_grad_net, loop_grad_map, loop_input_names, loop_output_names = \
        _gen_subnet_gradient(loop_net, init_grad_map)
    assert loop_grad_net, "Failed to get gradient net for loop body in While op"

    grad_ops += _prepare_gradient_while_ops(
        fwd_op=op,
        input_names=loop_input_names,
        output_names=loop_output_names,
        loop_grad_net=loop_grad_net,
        workspace_blob=workspace_blob,
        init_grad_map=init_grad_map,
        loop_grad_map=loop_grad_map)

    op_input = [str(i) for i in op.input]
    g_input = [loop_grad_map.get(i, None) for i in op_input]
    return grad_ops, g_input


# Constructs gradient While op, arguments:
#  fwd_op - forward While op
#  input_names - input blob names for a gradient op
#  output_names - output blob names for a gradient op
#  loop_grad_net - gradient loop body net
#  workspace_blob - blob that holds forward workspaces stack
#  init_grad_map - initial gradient to forward blob map
#  loop_grad_map - gradient blob map for loop's body
def _prepare_gradient_while_ops(
        fwd_op, input_names, output_names, loop_grad_net, workspace_blob,
        init_grad_map, loop_grad_map):
    gradient_while_def = caffe2_pb2.OperatorDef()
    gradient_while_def.CopyFrom(fwd_op)
    if gradient_while_def.name:
        gradient_while_def.name += "_grad"

    loop_net_arg = caffe2_pb2.Argument()
    loop_net_arg.name = "loop_net"
    loop_net_arg.n.CopyFrom(loop_grad_net)

    cond_net_arg = caffe2_pb2.Argument()
    cond_net_arg.name = "cond_net"
    from caffe2.python.core import Net, BlobReference
    # Construct condition net - check that there're still forward workspaces
    # left using HasScope op
    cond_net = Net('gradient_loop_cond_net')
    cond_init_net = Net('gradient_loop_cond_net_init')
    cond_blob = cond_net.NextScopedBlob(cond_net.Name() + '/cond')
    cond_init_net.HasScope(workspace_blob, cond_blob)
    cond_net.HasScope(workspace_blob, cond_blob)
    for blob, init_grad_blob in init_grad_map.items():
        blob_name = str(blob)
        init_grad_blob_name = str(init_grad_blob)
        if blob_name in loop_grad_map and \
                loop_grad_map[blob_name] != init_grad_blob_name:
            cond_net.Copy(
                BlobReference(loop_grad_map[blob_name]), init_grad_blob)
            cond_init_net.Copy(
                init_grad_blob, BlobReference(loop_grad_map[blob_name]))
    cond_net_arg.n.CopyFrom(cond_net.Proto())

    del gradient_while_def.arg[:]
    gradient_while_def.arg.extend([loop_net_arg, cond_net_arg])

    del gradient_while_def.control_input[:]
    del gradient_while_def.input[:]
    gradient_while_def.input.extend(
        [str(cond_blob).encode('utf-8')] + list(input_names))
    del gradient_while_def.output[:]
    gradient_while_def.output.extend(output_names)
    gradient_while_def.is_gradient_op = True
    return [o for o in cond_init_net.Proto().op] + [gradient_while_def]


def _get_do_arguments(do_op):
    assert do_op.type == "Do", "Expected Do op"
    args = {}
    for arg in do_op.arg:
        if not arg.name:
            continue
        if arg.name == "net":
            assert arg.n, "Expected non empty net argument"
            args["net"] = arg.n
        elif arg.name == "reuse_workspace":
            assert arg.i, "Expected non empty reuse_workspace argument"
            args["reuse_workspace"] = bool(arg.i)
        elif arg.name == "inner_blobs":
            assert arg.strings, "Expected non empty inner_blobs argument"
            args["inner_blobs"] = arg.strings
        elif arg.name == "outer_blobs_idx":
            assert arg.ints, "Expected non empty outer_blobs_idx argument"
            args["outer_blobs_idx"] = arg.ints
    return args


def gen_if_gradient(op, g_output):
    """
    Generates gradient If operator, given forward If op and a list
    of gradient blobs corresponding to forward op's outputs
    Returns a gradient op and a list of blobs corresponding to input gradients
    """
    from caffe2.python.core import BlobReference
    assert op.type == "If", "Expected If op"
    # first input is the condition blob
    assert len(op.input) > 0, "Expected at least one input in If op"

    assert len(op.output) == len(g_output), \
        "Different number of gradient blobs and If op outputs"

    grad_ops, deduped_g_output = dedupe_g_output(op, g_output)
    g_output = deduped_g_output

    init_grad_map = {}  # map from if's output blob to output gradient blob
    op_input = [str(i) for i in op.input]
    op_output = [str(o) for o in op.output]
    for output_name, grad_output_name in zip(op_output, g_output):
        if grad_output_name:
            init_grad_map[BlobReference(output_name)] = \
                BlobReference(grad_output_name)
    # shouldn't call without at least one output gradient available
    assert len(init_grad_map) > 0, "Empty initial gradient map for If op"

    grad_map = {}  # map from blob to gradient blob
    then_net = _get_net_argument(op, "then_net")
    assert then_net, "Expected then subnet in If op"
    then_grad_net, then_grad_map, then_input_names, then_output_names = \
        _gen_subnet_gradient(then_net, init_grad_map)
    assert then_grad_net, "Failed to get gradient net for then in If op"
    grad_map.update(then_grad_map)

    else_input_names = set()
    else_output_names = set()
    else_grad_map = {}
    else_grad_net = None
    else_net = _get_net_argument(op, "else_net")
    if else_net:
        else_grad_net, else_grad_map, else_input_names, else_output_names = \
            _gen_subnet_gradient(else_net, init_grad_map)
        assert else_grad_net, "Failed to get gradient net for else in If op"
        # consider case: else doesn't update blob's gradient and keeps original
        # from init_grad_map, but then updates the gradient
        for else_blob, else_grad_blob in else_grad_map.items():
            if else_blob in then_grad_map:
                then_grad_blob = then_grad_map[else_blob]
                # if both then and else branches have grad blob name for the same
                # blob and grad names are different, then one of the branches
                # doesn't use blob and has original grad blob name in it's grad map,
                # and another branch uses blob and has <blob_name>_grad name
                # in it's grad map (might be different from original grad blob)
                if then_grad_blob != else_grad_blob:
                    init_grad_name = init_grad_map[else_blob] \
                        if else_blob in init_grad_map else None

                    if then_grad_blob == init_grad_name:
                        grad_map[else_blob] = else_grad_blob
                    elif else_grad_blob == init_grad_name:
                        grad_map[else_blob] = then_grad_blob
                    else:
                        raise "Unexpected grad blob name " + else_blob + ", " + \
                            else_grad_blob + ", " + then_grad_blob
            else:
                grad_map[else_blob] = else_grad_blob

    # make sure gradients of blobs that were not computed
    # by the selected if's branch are initialized with zeros
    then_other_output_names = \
        then_output_names - (then_output_names & else_output_names)
    then_other_grad_output_names = set(
        [o for o in then_other_output_names if o in then_grad_map.values()])
    zero_then = _gen_grad_zero_init_ops(
        init_grad_map, then_grad_map, then_other_grad_output_names)
    if else_grad_net:
        else_grad_net.op.extend(zero_then)
    elif len(zero_then) > 0:
        else_grad_net = caffe2_pb2.NetDef()
        else_grad_net.CopyFrom(then_grad_net)
        if else_grad_net.name:
            else_grad_net.name += "_auto_else_zero_blobs_"
        del else_grad_net.op[:]
        else_grad_net.op.extend(zero_then)
        del else_grad_net.external_input[:]
        del else_grad_net.external_output[:]

    else_other_output_names = \
        else_output_names - (then_output_names & else_output_names)
    else_other_grad_output_names = set(
        [o for o in else_other_output_names if o in else_grad_map.values()])
    zero_else = _gen_grad_zero_init_ops(
        init_grad_map, else_grad_map, else_other_grad_output_names)
    then_grad_net.op.extend(zero_else)

    output_names = list(then_output_names | else_output_names)
    input_names = then_input_names | else_input_names
    # make sure condition blob is the first in the list
    input_names = [op_input[0]] + list(input_names - set(op_input[0]))
    gradient_if_def = _prepare_gradient_if_op(
        fwd_op=op,
        input_names=input_names,
        output_names=output_names,
        then_grad_net=then_grad_net,
        else_grad_net=else_grad_net)
    g_input = [grad_map.get(i, None) for i in op_input]
    return grad_ops + [gradient_if_def], g_input


def _gen_subnet_gradient(subnet, init_grad):
    grad_ops, grad_names_map = _gen_subgradient_pass(
        subnet, init_grad)

    output_names = set()
    input_names = set()
    for grad_op in grad_ops:
        for grad_op_input in grad_op.input:
            if str(grad_op_input) not in output_names:
                input_names.add(str(grad_op_input))
        for grad_op_output in grad_op.output:
            output_names.add(str(grad_op_output))

    gradient_net_def = caffe2_pb2.NetDef()
    gradient_net_def.CopyFrom(subnet)
    if gradient_net_def.name:
        gradient_net_def.name += "_grad"
    del gradient_net_def.op[:]
    gradient_net_def.op.extend(grad_ops)
    del gradient_net_def.external_input[:]
    del gradient_net_def.external_output[:]

    return gradient_net_def, grad_names_map, input_names, output_names


def _get_net_argument(op, net_name):
    for arg in op.arg:
        if arg.name and arg.name == net_name:
            assert arg.n, "Expected non empty net argument " + net_name
            return arg.n
    return None


def getNetArgument(op, net_name):
    """A wrapper for external call"""
    return _get_net_argument(op, net_name)


def _gen_subgradient_pass(subnet, init_grad):
    from caffe2.python.core import IR
    subnet_ir = IR(subnet.op)
    grad_ops, grad_blob_map = \
        subnet_ir.GetBackwardPass(init_grad)
    grad_names_map = {}
    for b, g in grad_blob_map.items():
        grad_names_map[str(b)] = str(g)
    return grad_ops, grad_names_map


def _do_op_sanity_check_and_process(op):
    assert op.type == "Do", "Expected Do op"

    subnet = _get_net_argument(op, "net")
    assert subnet, "No net argument found in Do op"

    inner_blobs = None
    outer_blobs_idx = None
    for arg in op.arg:
        if arg.name and arg.name == "inner_blobs":
            assert not inner_blobs, "inner_blobs redefinition"
            assert arg.strings and len(arg.strings) > 0, \
                "Empty inner_blobs argument in Do op"
            inner_blobs = [s.decode('utf-8') for s in arg.strings]
        if arg.name and arg.name == "outer_blobs_idx":
            assert not outer_blobs_idx, "outer_blobs_idx redefinition"
            assert arg.ints and len(arg.ints) > 0, \
                "Empty outer_blobs_idx argument in Do op"
            outer_blobs_idx = arg.ints
        if inner_blobs and outer_blobs_idx:
            break

    assert inner_blobs, "No inner_blobs argument found in Do op"
    assert outer_blobs_idx, "No outer_blobs_idx argument found in Do op"

    assert len(inner_blobs) == len(outer_blobs_idx), \
        "Arguments inner_blobs and outer_blobs_idx of different length in Do op"

    all_inner_blobs = set(inner_blobs)
    assert len(all_inner_blobs) == len(inner_blobs), \
        "Found duplicates in inner_blobs in Do op"

    op_input = [str(i) for i in op.input]
    assert len(op_input) > 0, "Expected at least one input blob"
    # remove last input blob that holds pointer to workspace
    input_workspace_blob_name = op_input[-1]
    op_input = op_input[:-1]

    op_output = [str(o) for o in op.output]
    assert len(op_output) > 0, "Expected at least one output blob"
    # remove last output blob that holds pointer to workspace
    workspace_blob_name = op_output[-1]
    assert input_workspace_blob_name == workspace_blob_name, \
        "Expected same input/output workspace blob"
    op_output = op_output[:-1]

    all_op_input_blob_names = set(op_input)
    assert len(all_op_input_blob_names) == len(op_input), \
        "Found duplicates in Do op inputs"
    all_op_output_blob_names = set(op_output)
    assert len(all_op_output_blob_names) == len(op_output), \
        "Found duplicates in Do op outputs"

    ordered_outer_blob_names = op_input + op_output
    all_outer_blob_names = set(ordered_outer_blob_names)
    used_outer_blob_names = set()
    outer_to_inner_map = {}
    inner_to_outer_map = {}
    for inner_name, outer_blob_idx in zip(inner_blobs, outer_blobs_idx):
        assert outer_blob_idx >= 0 and \
            outer_blob_idx < len(ordered_outer_blob_names), \
            "Outer blob index is out of bounds in Do op"
        outer_name = ordered_outer_blob_names[outer_blob_idx]
        assert outer_name not in used_outer_blob_names, \
            "Reusage of outer blob name " + outer_name + " in Do op"
        used_outer_blob_names.add(outer_name)
        outer_to_inner_map[outer_name] = inner_name
        inner_to_outer_map[inner_name] = outer_name

    assert len(used_outer_blob_names) == len(all_outer_blob_names), \
        "Not all outer blob names are used in blob bindings in Do op"

    return subnet, outer_to_inner_map, inner_to_outer_map, workspace_blob_name


def _prepare_blob_copy_op(from_name, to_name):
    copy_op_def = caffe2_pb2.OperatorDef()
    copy_op_def.type = "Copy"
    copy_op_def.input.extend([from_name])
    copy_op_def.output.extend([to_name])
    return copy_op_def


def _prepare_gradient_do_op(
        fwd_op, fwd_net, grad_ops, inputs, outputs, blob_bindings, saved_fwd_blobs,
        workspace_blob_name):
    gradient_net_def = caffe2_pb2.NetDef()
    gradient_net_def.CopyFrom(fwd_net)
    if gradient_net_def.name:
        gradient_net_def.name += "_grad"
    del gradient_net_def.op[:]
    gradient_net_def.op.extend(grad_ops)
    del gradient_net_def.external_input[:]
    del gradient_net_def.external_output[:]

    gradient_do_def = caffe2_pb2.OperatorDef()
    gradient_do_def.CopyFrom(fwd_op)
    if gradient_do_def.name and len(gradient_do_def.name) > 0:
        gradient_do_def.name += "_grad"

    del gradient_do_def.input[:]
    gradient_do_def.input.extend(inputs)
    # workspace pointer blob
    gradient_do_def.input.append(workspace_blob_name)
    del gradient_do_def.output[:]
    gradient_do_def.output.extend(outputs)
    # workspace pointer blob
    gradient_do_def.output.append(workspace_blob_name)

    net_arg = caffe2_pb2.Argument()
    net_arg.name = "net"
    net_arg.n.CopyFrom(gradient_net_def)

    ordered_new_outer_names = inputs + outputs
    inner_blobs = blob_bindings.keys()
    new_outer_blobs_idx = [ordered_new_outer_names.index(blob_bindings[b])
                            for b in inner_blobs]

    inner_blobs_arg = caffe2_pb2.Argument()
    inner_blobs_arg.name = "inner_blobs"
    inner_blobs_arg.strings.extend([b.encode('utf-8') for b in inner_blobs])

    outer_blobs_idx_arg = caffe2_pb2.Argument()
    outer_blobs_idx_arg.name = "outer_blobs_idx"
    outer_blobs_idx_arg.ints.extend(new_outer_blobs_idx)

    saved_blobs_arg = caffe2_pb2.Argument()
    saved_blobs_arg.name = "saved_fwd_blobs"
    saved_blobs_arg.strings.extend(
        [b.encode('utf-8') for b in saved_fwd_blobs])

    del gradient_do_def.arg[:]
    gradient_do_def.arg.extend([
        net_arg, inner_blobs_arg, outer_blobs_idx_arg, saved_blobs_arg])
    del gradient_do_def.control_input[:]

    gradient_do_def.is_gradient_op = True

    return gradient_do_def


def _gen_grad_zero_init_ops(init_grad_map, grad_map, grad_output_names):
    grad_init_ops = []
    for grad_output in grad_output_names:
        # get the corresponding output name blob and use it in ConstantFill
        # so that grad_output has the same shape
        output_name = None
        for o, g in grad_map.items():
            if g == grad_output:
                output_name = o
                break
        assert output_name, "Unknown gradient output " + grad_output

        grad_init_op = None
        # make sure that we do not overwrite existing gradients with zeros
        if output_name in init_grad_map:
            init_grad_name = init_grad_map[output_name]
            # in case we use a different gradient blob name, copy gradient
            if init_grad_name != grad_output:
                grad_init_op = caffe2_pb2.OperatorDef()
                grad_init_op.type = "Copy"
                grad_init_op.input.extend([str(init_grad_name)])
                grad_init_op.output.extend([str(grad_output)])
        else:
            grad_init_op = caffe2_pb2.OperatorDef()
            grad_init_op.type = "ConstantFill"
            grad_init_op.input.extend([output_name])
            grad_init_op.output.extend([grad_output])
            value_arg = caffe2_pb2.Argument()
            value_arg.name = "value"
            value_arg.f = 0.0
            grad_init_op.arg.extend([value_arg])

        if grad_init_op:
            grad_init_ops.append(grad_init_op)
    return grad_init_ops


def _prepare_gradient_if_op(
        fwd_op, input_names, output_names, then_grad_net, else_grad_net):
    gradient_if_def = caffe2_pb2.OperatorDef()
    gradient_if_def.CopyFrom(fwd_op)
    del gradient_if_def.input[:]
    gradient_if_def.input.extend(input_names)
    del gradient_if_def.output[:]
    gradient_if_def.output.extend(output_names)

    then_net_arg = caffe2_pb2.Argument()
    then_net_arg.name = "then_net"
    then_net_arg.n.CopyFrom(then_grad_net)
    gradient_args = [then_net_arg]
    if else_grad_net:
        else_net_arg = caffe2_pb2.Argument()
        else_net_arg.name = "else_net"
        else_net_arg.n.CopyFrom(else_grad_net)
        gradient_args.append(else_net_arg)

    del gradient_if_def.arg[:]
    gradient_if_def.arg.extend(gradient_args)
    if gradient_if_def.name:
        gradient_if_def.name += "_grad"
    del gradient_if_def.control_input[:]
    gradient_if_def.is_gradient_op = True
    return gradient_if_def


def disambiguate_grad_if_op_output(grad_op, idx, new_grad_output):
    then_net = _get_net_argument(grad_op, "then_net")
    old_grad_out_match = grad_op.output[idx]
    for op in then_net.op:
        for i, out in enumerate(op.output):
            if out == old_grad_out_match:
                op.output[i] = new_grad_output
    else_net = _get_net_argument(grad_op, "else_net")
    if else_net:
        for op in else_net.op:
            for i, out in enumerate(op.output):
                if out == old_grad_out_match:
                    op.output[i] = new_grad_output
    grad_op.output[idx] = new_grad_output