code: def test_mlp_learnt_activation():
      mlp = eqx.nn.MLP(
          2,
          5,
          8,
          2,
          activation=eqx.nn.PReLU(),
          final_activation=eqx.nn.PReLU(),
          key=jrandom.PRNGKey(5678),
      )
      x = jnp.array([0.5, 0.7])
      assert mlp.activation.negative_slope.shape == (2, 8)
      assert mlp.final_activation.negative_slope.shape == (5,)
  
      @eqx.filter_jit
      @eqx.filter_grad
      def grad(mlp, x):
          return jnp.sum(mlp(x))
  
      grads = grad(mlp, x)
      assert grads.activation.negative_slope.shape == (2, 8)
      assert grads.final_activation.negative_slope.shape == (5,)
  
imports: import warnings
  from typing import Union
  import equinox
  import jax
  import jax.nn
  import jax.numpy
  import jax.random
  import pytest
label_desc: Use image node_id 12 for calling a global function with example usage: test_mlp_learnt_activation() and returns: jnp