How to separate the Jacobian based on effectively?

[mathliuqiang]

I have constructed an MLP network function with a structure of [2, 8, 2], meaning the output dimension is 2.

import jax
from jax import random
import jax.numpy as jnp

# ------ Code from Jax documentation -----
def random_layer_params(m, n, key, scale=1e-1):
  w_key, b_key = random.split(key)
  return scale * random.normal(w_key, (n, m)), scale * random.normal(b_key, (n,))

def init_params(sizes, key):
  keys = random.split(key, len(sizes))
  return [random_layer_params(m, n, k) for m, n, k in zip(sizes[:-1], sizes[1:], keys)]

# mlp is network
def mlp(params, inpt):
    hidden = inpt
    for w, b in params[:-1]:
        outputs = jnp.dot(w, hidden) + b
        hidden = jax.nn.sigmoid(outputs)

    final_w, final_b = params[-1]
    hidden = jnp.dot(final_w, hidden) + final_b
    return hidden

I am using jax.jacobian to compute the derivative with respect to the parameters.

layers = [2, 8, 2]
params = init_params(layers, jax.random.key(0))

x = jnp.array([1.0, 3.0])
y = mlp(params, x)
dy_fun = jax.jacobian(mlp)
dy = dy_fun(params, x)

print(f'The value of network is : {y}')

print(f'The shape of Jacobi matrix as follows:')
print(dy[0][0].shape)
print(dy[0][1].shape)
print(dy[1][0].shape)
print(dy[1][1].shape)

The output is

The value of network is : [-0.16414015  0.0424983 ]
The shape of Jacobi matrix as follows:
(2, 8, 2)
(2, 8)
(2, 2, 8)
(2, 2)

However, I want to separate the resulting Jacobian based on the output dimension. Specifically, I need to extract the derivatives for each of the two outputs individually, i.e., $\frac{\partial \mathop{output}_1}{\partial \theta_i}$ and $\frac{\partial \mathop{output}_2}{\partial \theta_i}$, ideally representing them as two separate Pytrees (or vectors), where each Pytree contains the gradients for one output element with respect to all parameters. How can I achieve this?

[jakevdp]

With params being a nested list of tuples in your example, can you be more specific about what you mean by $\theta_i$? What output shapes are you hoping to see?