Unclear how to add up two grads of nn.Module type

[BoyuanChen99]

Description of the model to be implemented

I am trying to add up the grads for an optimizer. The grads has type flax.nn.Module. I wonder how to sum them up.

class NeRF_Model(nn.Module):
    def apply(self, x):
        the_input = x
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = jnp.append(arr=x, values=the_input, axis=1)  # Dense_4: (257, 256)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=4, kernel_init = initializers.normal(1),
                     bias_init=initializers.ones)
        # output: (r,g,b, \upsigma)
        return x

Dataset the model could be trained on

Image Data

Specific points to consider

/

Reference implementations in other frameworks

/

[levskaya]

I believe you mean the grads are stored inside an nn.Model i.e the optimizer.target?
so if you do

_, params = NeRF_Model.init(random.PRNGKey(0), jnp.ones((10,10)))
model = nn.Model(NeRF_Model, params)
optimizer = flax.optim.Adam(learning_rate=0.01).create(model)
del model
# optimizer.target contains the model

you could then use a tree_multimap:

def dumb_loss(model, x):
  return jnp.sum(model(x))
grad_fn = jax.grad(dumb_loss)
x1 = random.uniform(random.PRNGKey(0), (10, 10))
x2 = random.uniform(random.PRNGKey(1), (10, 10))
grad1 = grad_fn(optimizer.target, x1)
grad2 = grad_fn(optimizer.target, x2)
summed_grad = jax.tree_multimap(lambda x,y: x+y, grad1, grad2)

in a colab, see:
https://colab.research.google.com/drive/13UCy6kLRO8u3NP_lVQl9rhh_3_FzhB72?usp=sharing