Define jacobian with tensorflow custom_gradient

Question

I'm trying to define a custom gradient in TensorFlow for a function that returns a non-scalar.

A simplified toy version of the problem is:

import tensorflow as tf

@tf.custom_gradient
def f(x):
    y = tf.Variable([x[1], 2.0*x[0]])

    def grad(upstream):
        dg = tf.Variable([[0.0, 2.0], [1.0, 0.0]])
        return dg * upstream # ? no idea here

    return y, grad

x = tf.Variable([1.0, 2.0])
with tf.GradientTape() as tape:
    y = f(x)
tape.jacobian(y, x)
# tape.gradient(y,x)

I do not understand what shape the grad function needs to return. In the above example, tape.gradient gives the full jacobian of the problem (which I would expect from tape.jacobian)and tape.jacobian gives nothing (though I had been getting error messages before restarting my kernel).

Is this the expected shape and behaviour? What shape should the custom gradient return for a vector-valued or matrix-valued function's derivative?

user_4279607 · Answer

I agree it can be confusing. Let me try to expand the statement given in the Tensorflow docs and draw an explicit link between the maths and the code.

Let us consider a function

$f:\mathbb{R}^m ightarrow \mathbb{R}^n$ .

Its Jacobian is the n x m matrix:

$\frac{\partial f}{\partial x} = \begin{bmatrix} \frac{\partial f_1}{\partial x_1}&\cdots&\frac{\partial f_1}{\partial x_m}\vdots&\ddots&\vdots\frac{\partial f_n}{\partial x_1}&\cdots&\frac{\partial f_n}{\partial x_m} \end{bmatrix}$ .

Given an input vector x and an "upstream" vector u, the function grad_fn in a @custom_gradient function should return the "vector-Jacobian product" vjp, i.e.:

$\mathrm{vjp}(f)(u, x) = u^ op \frac{\partial f}{\partial x}(x)$ .

Note that in this formulation u is given and does not depend on x. Therefore, we can rewrite the vjp as follows:

$\mathrm{vjp}(f)(u, x) = \frac{\partial}{\partial x}((u^ op f(x))$ .

This can be generalised to higher-rank tensors if we define $<a,b>$ as the sum of the elementwise product of a and b (i.e. tf.reduce_sum(a * b)), which is the regular dot product for vectors and the Frobenius inner product for matrices. Then, the output of grad_fn should be a tensor with the same shape as x computed as:

$\Big(\mathrm{vjp}(f)(u, x)\Big){i,j,\dots} = \frac{\partial}{\partial x{i,j,\dots}}\langle u,f(x) angle$ .

Note that this is effectively the gradient of the scalar-valued function tf.reduce_sum(u * f(x)) (and is therefore easier to think about and to work with).

In your specific case, you need to take the gradient of u[0]*x[1]+u[1]*2.0*x[0] with respect to x[0],x[1], i.e. something like:

def grad(upstream):
    return tf.Variable([2.0*upstream[1], upstream[0]])

This produces the desired gradient.

Unfortunately jacobian won't work with this (as it doesn't with your original code) because your code has a second issue: by using variables it breaks the link between the outputs and the inputs (creating a variable is like copy/pasting the values and does not propagate the gradients). You can fix this as follows:

@tf.custom_gradient
def f(x):
    y = tf.stack([x[1], 2.0*x[0]])

    def grad(upstream):
        return tf.stack([2.0*upstream[1], upstream[0]])

    return y, grad

Define jacobian with tensorflow custom_gradient

1 Answers