Skip to content

Add option to disable OpFromGraph.L_Op/R_Op from being wrapped in their own OpFromGraph #1174

Open
@ricardoV94

Description

@ricardoV94

The fact that we wrap a graph in OpFromGraph, doesn't mean we also want to wrap it's gradient. It can also lead to troubles as in pymc-devs/pymc#7657

I suggest adding a wrap_grad_in_ofg: bool kwarg that toggles this behavior. The easiest way to get this change is to keep the behavior as is, but call the inline_ofg rewrite manually after calling it here:

connected_input_grads = iter(
lop_op(*inputs, *connected_outputs, *connected_output_grads, **kwargs)
)

And here:

connected_output_grads = iter(rop_op(*inputs, **kwargs))

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions