Hessian/Laplacian of an already batched neural network

[Binbose]

I have a neural network function that maps (batch_size, x, n_help_particle, n_space_dim) -> (batch_size,) with x having the shape n_particle, n_space_dim and p having the shape n_particle_2, n_space_dim:

class SimpleNet(nn.Module):
    num_layers: int
    num_hidden: int
    activation: nn.tanh
    vector_output: bool
    dtype: object

    def setup(self) -> None:
        self.simple_mlp1 = nn.Dense(features=self.num_hidden,
                                    dtype=self.dtype,
                                    param_dtype=self.dtype,
                                    kernel_init=nn.initializers.lecun_normal(dtype=self.dtype),
                                    bias_init=nn.initializers.normal(dtype=self.dtype))
        self.simple_mlp2 = nn.Dense(features=self.num_hidden,
                                    dtype=self.dtype,
                                    param_dtype=self.dtype,
                                    kernel_init=nn.initializers.lecun_normal(dtype=self.dtype),
                                    bias_init=nn.initializers.normal(dtype=self.dtype))
        self.simple_mlp3 = nn.Dense(features=2,
                                    dtype=self.dtype,
                                    param_dtype=self.dtype,
                                    kernel_init=nn.initializers.lecun_normal(dtype=self.dtype),
                                    bias_init=nn.initializers.normal(dtype=self.dtype))

    def __call__(self, x, potential, *args, **kwargs):
        n_particle_original = x.shape[1]
        x = jnp.concatenate([x, potential], axis=1)

        y = self.simple_mlp1(x)
        y = nn.tanh(y)
        y = self.simple_mlp2(y)
        y = nn.tanh(y)
        y = self.simple_mlp3(y)

        return ((y[:,:n_particle_original] - x[:,:n_particle_original])**2).sum(-1)

rng_key = jax.random.PRNGKey(0)
rng, input_rng, init_rng = jax.random.split(rng_key, 3)
batch_size = 4
n_potential = 1
n_particle = 1
n_space_dim = 2 #1
vector_output = True
input = jax.random.uniform(input_rng, shape=(batch_size, n_particle, n_space_dim))
potential = jax.random.uniform(init_rng, shape=(batch_size, n_potential, n_space_dim))

net = SimpleNet(num_layers=2, num_hidden=32, activation=nn.elu, vector_output=vector_output, dtype=jnp.float32)
params = net.init(rng, input, potential)
netf = lambda x: net.apply(params, x, potential).sum(-1)
netf_h = jax.hessian(netf)
netf_h(input)

I want to take the hessian of the function with respect to the x input, so I would expect something with the output (batch_size, n_particle, n_space_dim, n_particle, n_space_dim), but I am always getting something of the shape (batch_size, batch_size, n_particle, n_space_dim, batch_size, n_particle, n_space_dim). I assume the reason for this is that my function already takes in a batched input, so one remedy would be to rewrite everything without a batch dimension, but is there an easier way that can calculate the hessian (and Laplacian) of an already batched function?

[jakevdp]

If you have a function that truly operates independently across batches, you can do this by wrapping your jax.hessian call in a vmap. In your case it might look like this:

netf = lambda x, p: net.apply(params, x, p).sum(-1)
out = jax.vmap(jax.hessian(netf, argnums=0))(input, potential)
print(out.shape)
# (4, 1, 2, 1, 2)

Note that because both input and potential are batched, they both need to be arguments to the vmap-transformed function.
Additionally, in order for this to work, your code needs to be written in a way that correctly handles un-batched inputs. It looks like in your case that can be done by changing these lines:

        n_particle_original = x.shape[1]
        x = jnp.concatenate([x, potential], axis=1)

to this:

        n_particle_original = x.shape[-2]
        x = jnp.concatenate([x, potential], axis=-2)

[Binbose]

Awesome, thank you!