code: def big_resnet(
      scope: Scope,
      x,
      blocks=(10, 5),
      dtype=jnp.float32,
      norm=default_norm,
      act=nn.relu,
  ):
      conv = partial(nn.conv, bias=False, dtype=dtype)
      norm = partial(norm, dtype=dtype)
  
      # a two stage resnet where inner blocks are rematerialized to make sure
      # memory consumtion grows as O(sqrt(N)) and compute is O(N) where N is the number of blocks..
      # we use a double scan such that the compiled binary is of size O(1).
      print("total residual blocks:", np.prod(blocks))
  
      def body_fn(scope, x):
          return residual_block(
              scope, x, conv, norm, act, features=x.shape[-1]
          )
  
      return lift.remat_scan(
          body_fn,
          lengths=blocks,
          variable_axes={"params": 0, "batch_stats": 0},
          split_rngs={"params": True},
          policy=None,
      )(scope, x)
  
file_path: flax/tests/core/design/core_big_resnets_test.py
imports: from functools import partial
  import jax
  import numpy
  from absl.testing import absltest
  from jax import numpy
  from jax import random
  from flax.core import Array, Scope, init, lift, nn, unfreeze
label_desc: Use image node_id 2 for calling a global function with example usage: big_resnet(scope, x, blocks, dtype, norm, act) and returns: lift, residual_block