code: def transpose_last_axis_to_first(
      tree: chex.ArrayTree,
  ) -> chex.ArrayTree:
      """Function to transpose the last axis to be first for all leaves in a pytree.
  
      This function will transpose the last axis to the front for all leaves in a
      pytree; each leaf with shape [D_1, ..., D_{n-1}, D_n] will have shape
      [D_n, D_1, ..., D_{n-1}].
  
      Args:
        tree: the pytree of Arrays to be transposed.
  
      Returns:
        tree: the transposed output tree.
      """
  
      def _transpose(tree):
          return jnp.transpose(
              tree, [tree.ndim - 1] + list(range(tree.ndim - 1))
          )
  
      return tree_fn(_transpose)(tree)
  
file_path: rlax/rlax/_src/tree_util.py
imports: import functools
  from typing import Any, Callable, Sequence
  import chex
  import jax
  import jax.numpy
  import numpy
  from rlax._src import base
label_desc: Use image node_id 7 for calling a global function with example usage: transpose_last_axis_to_first(tree) and returns: tree_fn, jnp