L operator and R operator in Tensorflow

grad_lib.py

import tensorflow as tf


def gradients(f, x, grad_ys=None):
    '''
    An easier way of computing gradients in tensorflow. The difference from tf.gradients is
        * If f is not connected with x in the graph, it will output 0s instead of Nones. This will be more meaningful
            for computing higher-order gradients.

        * The output will have the same shape and type as x. If x is a list, it will be a list. If x is a Tensor, it
            will be a tensor as well.

    :param f: A `Tensor` or a list of tensors to be differentiated
    :param x: A `Tensor` or a list of tensors to be used for differentiation
    :param grad_ys: Optional. It is a `Tensor` or a list of tensors having exactly the same shape and type as `f` and
                    holds gradients computed for each of `f`.
    :return: A `Tensor` or a list of tensors having the same shape and type as `x`
    '''

    if isinstance(x, list):
        grad = tf.gradients(f, x, grad_ys=grad_ys)
        for i in range(len(x)):
            if grad[i] is None:
                grad[i] = tf.zeros_like(x[i])
        return grad
    else:
        grad = tf.gradients(f, x, grad_ys=grad_ys)[0]
        if grad is None:
            return tf.zeros_like(x)
        else:
            return grad


def Lop(f, x, v):
    '''
    Compute Jacobian-vector product. The result is v^T @ J_x

    :param f: A `Tensor` or a list of tensors for computing the Jacobian J_x
    :param x: A `Tensor` or a list of tensors with respect to which the Jacobian is computed.
    :param v: A `Tensor` or a list of tensors having the same shape and type as `f`
    :return: A `Tensor` or a list of tensors having the same shape and type as `x`
    '''
    assert not isinstance(f, list) or isinstance(v, list), "f and v should be of the same type"
    return gradients(f, x, grad_ys=v)


def Rop(f, x, v):
    '''
    Compute Jacobian-vector product. The result is J_x @ v.
    The method is inspired by [deep yearning's blog](https://j-towns.github.io/2017/06/12/A-new-trick.html)
    :param f: A `Tensor` or a list of tensors for computing the Jacobian J_x
    :param x: A `Tensor` or a list of tensors with respect to which the Jacobian is computed
    :param v: A `Tensor` or a list of tensors having the same shape and type as `v`
    :return: A `Tensor` or a list of tensors having the same shape and type as `f`
    '''
    assert not isinstance(x, list) or isinstance(v, list), "x and v should be of the same type"
    if isinstance(f, list):
        w = [tf.ones_like(_) for _ in f]
    else:
        w = tf.ones_like(f)
    return gradients(Lop(f, x, w), w, grad_ys=v)