Seperating equations using Jaxpr

[pranavps-creator]

Hello Everyone,

First of all, I would like to apologize for the long post. I am working on a Physics-Informed Neural Networks project using an FNN-based Architecture called DeepONet. The physics part(also known as residual net in the code) is a function of the network output. I want to develop a function that identifies the eqns in the Jaxpr object that obtained from the Neural Network. Below is a small reproducible example of the problem:

import jax.random as random 
import jax.numpy as jnp
from jax import grad
from jax import vmap
import jax
from jax.nn import relu 
#Define MLP
def MLP(layers, activation=relu):
    ''' Vanilla MLP'''

    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

branch_params = [1,200,200]
trunk_params = [1,200,200]

#Define ANN functions
branch_init, branch_apply = MLP(branch_params)
trunk_init, trunk_apply= MLP(trunk_params)

branch_key, trunk_key = random.split(random.PRNGKey(1),2)
# Initialize weights
branch_params= branch_init(branch_key)
trunk_params= trunk_init(trunk_key)

#Obtain network output
def Operator_Net(branch_params, trunk_params, branch_input, trunk_input):
    
    #Branch output
    B = branch_apply(branch_params, branch_input)
    #Trunk output
    T = trunk_apply(trunk_params,trunk_input)
    
    out1 = jnp.sum(B[:100]*T[:100])
    out2 = jnp.sum(B[100:]*T[100:])
    
    return out1, out2

def out1(branch_params, trunk_params, branch_input, trunk_input):
    
    out1,_ = Operator_Net(branch_params, trunk_params, branch_input, trunk_input)
    return out1
    
    
def out2(branch_params, trunk_params, branch_input, trunk_input):
    _,out2 = Operator_Net(branch_params, trunk_params, branch_input, trunk_input)
    return out2

# Physics applied on the network output  
def residual_net(branch_params, trunk_params, branch_input, trunk_input):
    res1 = vmap(grad(out1,argnums=2),(None,None,0,0))(branch_params, trunk_params, branch_input, trunk_input)+vmap(grad(out2,argnums=2),(None,None,0,0))(branch_params, trunk_params, branch_input, trunk_input)
    res2 = vmap(grad(out1,argnums=3),(None,None,0,0))(branch_params, trunk_params, branch_input, trunk_input)+vmap(grad(out2,argnums=3),(None,None,0,0))(branch_params, trunk_params, branch_input, trunk_input)
    return res1,res2


#Define the jax_pr 
jaxpr = jax.make_jaxpr(residual_net)(branch_params,trunk_params,jnp.ones(1),jnp.ones(1))
print(jaxpr)

From this jaxpr, how do I identify the eqns that belong to the 'Operator_Net' function or in other words how do I identify the equations that build-up to the output of the 'Operator_Net'?

Thanks in advance