2nd order derivative is returning a Traced array with DynamicJaxprTrace instead of Traced array with BatchTrace

[pranavps-creator]

I am working on a Physics-Informed Neural Network project using Jax, where I want to compute the second-order derivative of the output of a Neural network with respect to one of its derivatives(x). The code goes as below:

def residual_net(self, params, E,E_physics,x, y):
    ux_x =grad(self.ux_fn, argnums = 2)(params, E, x, y)
    ux_xx  =grad(grad(self.ux_fn, argnums = 2), argnums = 2)(params, E, x, y)
    return ux_x, ux_xx

Here, the function ux_fn returns the output of the networks. When I compute the ux_x, it returns a Traced<ShapedArray(float32[])>with<BatchTrace(level=4/0)> with reasonable values. However, ux_xx returns a Traced<ShapedArray(float32[])>with<DynamicJaxprTrace(level=1/0)> with no value. How can I avoid this and compute the actual second-order derivative, i.e. d²u/dx²?

Note: The function residual_net is called using a vmap as:
vmap(self.residual_net, (None, 0,0, 0, 0))(params, E_branch, E_physics,coords[:, 0], coords[:, 1])

to compute the losses.

[jakevdp]

The various tracer classes used are internal implementation details, and users shouldn't have to worry about the details.

Can you add more detail to your question, in particular, how you're encountering the traced values? (a minimal reproducible example would be helpful here). If you're simply executing transformed functions, the inputs and outputs should be normal JAX arrays. Tracers are only used internally.

[pranavps-creator]

Hi @jakevdp, Thank you for your response. The problem lies with calculating the second-order derivatives. I always get zeros although I expected non-zero values. Here's a minimal reproducible example:

import numpy as np
import jax.numpy as jnp
from jax import random, grad, vmap, jit,jacfwd,jacrev
from jax.example_libraries import optimizers
from jax.nn import relu

#Define the Multi-Layer perception
def MLP(layers, activation=relu):
   

    def init(rng_key):
        def init_layer(key, d_in, d_out):
            k1, k2 = random.split(key)
            glorot_stddev = 1. / jnp.sqrt((d_in + d_out) / 2.)
            W = glorot_stddev * random.normal(k1, (d_in, d_out))
            b = jnp.zeros(d_out)
            return W, b

        key, *keys = random.split(rng_key, len(layers))
        params = list(map(init_layer, keys, layers[:-1], layers[1:]))
        return params

    def apply(params, inputs):
        for W, b in params[:-1]:
            outputs = jnp.dot(inputs, W) + b
            inputs = activation(outputs)
        W, b = params[-1]
        outputs = jnp.dot(inputs, W) + b
        return outputs

    return init, apply

layers = [2, 400, 400, 400, 400, 400, 400, 400, 400]

init, apply = MLP(
                layers, activation=relu
            )

x = jnp.arange(0,1,0.0025)
y = jnp.arange(0,1,0.0025)

params = init(random.PRNGKey(1))

def Operator_Net(params,x,y):
    z=jnp.stack([x,y])
    z=jnp.transpose(z)
    u= jnp.sum(apply(params,z))
    return u

def gradients(params,x,y):
    u_x = vmap(grad(Operator_Net, argnums =  1),(None,0,0))(params,  x, y)
    return u_x

def loss(params,x,y):
    u = vmap(Operator_Net,(None,0,0))(params,x,y)
    u_x = grad(Operator_Net,argnums=1)(params,x,y)
    u_xx = vmap(grad(grad(Operator_Net,argnums=1),argnums=1),(None,0,0))(params,x,y)
    return u_x,u_xx

u_x,u_xx= loss(params,x,y)
print(u_x)
print(u_xx)

[jakevdp]

I think the second derivative is expected to be zero in this case, because your operation is a sequence of matmuls and relus: matmul is linear (so its second derivative is zero) and relu has a constant gradient everywhere that the gradient is defined (so its second derivative is zero everywhere).

A sequence of operations with zero second derivative will result in a function with zero second derivative.

[pranavps-creator]

Thanks, @jakevdp . I forgot the role that the activation function played. B<ý changing the activation function to Tanh (jnp.tanh), I was able to get non zero values.

[jakevdp]

Yeah, that sounds right. Glad the issue is resolved for you!