Low-level graph transformation that transforms a graph into its gradient.
The function f represented by graph must return a single
float scalar. The resulting graph computes the gradients of that scalar with respect
to the inputs specified by wrt.
Arguments
- graph
(
AnvlGraph)
The graph to transform. Must produce a single scalar float output.- wrt
(
NULL|character())
Names of the graph inputs to differentiate with respect to.NULLdifferentiates with respect to all inputs.
Value
(AnvlGraph)
Its outputs are the requested gradients.
Details
To support alternative forward passes for more efficient backward passes, we replay and possibly rewrite the graph into a new descriptor. Afterwards, we traverse it backwards and call the gradient rules where necessary.
See rule_reverse() for more information.
gradient() and value_and_gradient() differentiate the same way, but
into the trace they are called in rather than into a graph of its own;
prefer them unless you need to operate on graphs directly.
Examples
graph <- trace_fn(prim_mul, list(nv_aval("f32", integer()), nv_aval("f32", integer())))
graph
#> <AnvlGraph> (%x1: f32[], %x2: f32[]) {
#> %1: f32[] = mul(%x1, %x2)
#> return %1
#> }
graph_gradient(graph, "lhs")
#> <AnvlGraph> [%c1: f32[]] (%x1: f32[], %x2: f32[]) {
#> %1: f32[] = mul(%x1, %x2)
#> %2: f32[] = mul(%c1, %x2)
#> return %2
#> }