test_higher_order_grad.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.


import math
from mxnet import nd, autograd
from mxnet.test_utils import assert_almost_equal, random_arrays, rand_shape_nd
from common import with_seed


@with_seed()
def test_sin():
    def sin(x):
        return nd.sin(x)

    def grad_grad_op(x):
        return -nd.sin(x)

    for dim in range(1, 5):
        shape = rand_shape_nd(dim)
        array = random_arrays(shape)
        check_second_order_unary(array, sin, grad_grad_op)


@with_seed()
def test_cos():
    def cos(x):
        return nd.cos(x)

    def grad_grad_op(x):
        return -nd.cos(x)

    for dim in range(1, 5):
        shape = rand_shape_nd(dim)
        array = random_arrays(shape)
        check_second_order_unary(array, cos, grad_grad_op)


@with_seed()
def test_relu():
    def relu(x):
        return nd.relu(x)

    def grad_grad_op(x):
        return nd.zeros_like(x)

    for dim in range(1, 5):
        shape = rand_shape_nd(dim)
        array = random_arrays(shape)
        check_second_order_unary(array, relu, grad_grad_op)


@with_seed()
def test_log():
    def log(x):
        return nd.log(x)

    def grad_grad_op(x):
        return -1/(x**2)

    for dim in range(1, 5):
        shape = rand_shape_nd(dim)
        array = random_arrays(shape)
        check_second_order_unary(array, log, grad_grad_op)


@with_seed()
def test_log2():
    def log2(x):
        return nd.log2(x)

    def grad_grad_op(x):
        return -1/((x**2) * math.log(2))

    for dim in range(1, 5):
        shape = rand_shape_nd(dim)
        array = random_arrays(shape)
        check_second_order_unary(array, log2, grad_grad_op)


@with_seed()
def test_log10():
    def log10(x):
        return nd.log10(x)

    def grad_grad_op(x):
        return -1/((x**2) * math.log(10))

    for dim in range(1, 5):
        shape = rand_shape_nd(dim)
        array = random_arrays(shape)
        check_second_order_unary(array, log10, grad_grad_op)


@with_seed()
def test_sigmoid():
    def sigmoid(x):
        return nd.sigmoid(x)

    def grad_grad_op(x):
        # Actual: f(x) * f'(x)
        # Expected: f(x)
        return sigmoid(x) * (1 - sigmoid(x)) * sigmoid(x)

    for dim in range(1, 5):
        shape = rand_shape_nd(dim)
        array = random_arrays(shape)
        check_sigmoid(array, sigmoid, grad_grad_op)


def check_second_order_unary(x, op, grad_grad_op):
    x = nd.array(x)
    expect_grad_grad = grad_grad_op(x)
    x.attach_grad()
    with autograd.record():
        y = op(x)
        y_grad = autograd.grad(y, x, create_graph=True, retain_graph=True)[0]
    y_grad.backward()
    assert_almost_equal(expect_grad_grad.asnumpy(), x.grad.asnumpy())


def check_sigmoid(x, op, grad_grad_op):
    x = nd.array(x)
    grad_grad_x = grad_grad_op(x)
    x.attach_grad()

    y_grad = nd.ones_like(x) *  nd.random.normal(shape=x.shape) 
    head_grad_grad = nd.ones_like(x) * nd.random.normal(shape=x.shape) 

    with autograd.record():
        y = op(x)
        x_grad = autograd.grad(y, x, y_grad,create_graph=True, retain_graph=True)[0]
        # x_grad.backward(head_grad_grad)
        x_grad_grad = autograd.grad(x_grad, [x], head_grad_grad,create_graph=False, retain_graph=True)[0]
    expected_grad_grad = grad_grad_x * head_grad_grad # * y_grad
    assert_almost_equal(expected_grad_grad.asnumpy(), x_grad_grad.asnumpy())
    # assert_almost_equal(expected_grad_grad.asnumpy(), x.grad.asnumpy())


if __name__ == '__main__':
    import nose
    nose.runmodule()