code: def test_transpose_last_axis_to_first(self):
      tree = {
          "a": jnp.ones((1, 2, 3, 4)),
          "b": {"c": jnp.ones((1, 2, 3, 4))},
      }
      expected_tree = {
          "a": jnp.ones((4, 1, 2, 3)),
          "b": {"c": jnp.ones((4, 1, 2, 3))},
      }
      chex.assert_trees_all_equal_shapes(
          tree_util.transpose_last_axis_to_first(tree["a"]),
          expected_tree["a"],
      )
      chex.assert_trees_all_equal_shapes(
          tree_util.transpose_last_axis_to_first(tree), expected_tree
      )
  
file_path: rlax/rlax/_src/tree_util_test.py
imports: from absl.testing import absltest
  import chex
  import jax
  import jax.numpy
  import numpy
  from rlax._src import tree_util
label_desc: Use image node_id 11 for calling the TreeUtilTest obj's underlying member method code with example usage: obj.test_transpose_last_axis_to_first() without return types