code: def transpose(fn):
      def trans(variables):
          return jax.tree_util.tree_map(lambda x: x.T, variables)
  
      return lift.map_variables(
          fn, "params", map_in_fn=trans, map_out_fn=trans, mutable=True
      )
  
file_path: flax/tests/core/design/core_tied_autoencoder_test.py
imports: from dataclasses import dataclass
  import jax
  from absl.testing import absltest
  from jax import numpy
  from jax import random
  from flax.core import init, lift, nn, unfreeze
label_desc: Use image node_id 1 for calling a global function with example usage: transpose(fn) and returns: lift, jax