code: def mlp_custom_grad(
      scope: Scope,
      x: Array,
      sizes: Sequence[int] = (8, 1),
      act_fn: Callable[[Array], Array] = nn.relu,
  ):
      f = nn.dense
  
      def fwd(scope, x, features):
          y, vjp_fn = lift.vjp(partial(f, features=features), scope, x)
          return y, vjp_fn
  
      def bwd(features, res, y_t):
          del features
          vjp_fn = res
          params_t, *input_t = vjp_fn(y_t)
          params_t = jax.tree_util.tree_map(jnp.sign, params_t)
          return (params_t, *input_t)
  
      dense_custom_grad = lift.custom_vjp(
          f, forward_fn=fwd, backward_fn=bwd, nondiff_argnums=(2,)
      )
  
      # hidden layers
      for size in sizes[:-1]:
          x = scope.child(dense_custom_grad, prefix="hidden_")(x, size)
          x = act_fn(x)
  
      # output layer
      return scope.child(dense_custom_grad, "out")(x, sizes[-1])
  
file_path: flax/tests/core/design/core_custom_vjp_test.py
imports: from functools import partial
  from typing import Callable, Sequence
  import jax
  import numpy
  from absl.testing import absltest
  from jax import numpy
  from jax import random
  from flax.core import Array, Scope, apply, init, lift, nn, unfreeze
label_desc: Use image node_id 1 for calling a global function with example usage: mlp_custom_grad(scope, x, sizes, act_fn) and returns: scope, y,  vjp_fn, params_t,  unknown