code: def test_tree_split_leaves(self):
      t = {
          "a0": np.zeros(3),
          "d": {
              "a1": np.arange(3),
              "a2": np.zeros([3, 3]) + 2,
          },
          "t4": (np.zeros([3, 2]) + 4, np.arange(3)),
      }
  
      for keepdim in (False, True):
          expd_shapes = jax.tree_map(lambda x: np.zeros(x.shape[1:]), t)
          if keepdim:
              expd_shapes = jax.tree_map(
                  lambda x: np.expand_dims(x, 0), expd_shapes
              )
  
          res_trees = tree_util.tree_split_leaves(
              t, axis=0, keepdim=keepdim
          )
          self.assertLen(res_trees, 3)
          chex.assert_trees_all_equal_shapes(expd_shapes, *res_trees)
          for i, res_t in enumerate(res_trees):
              np.testing.assert_allclose(res_t["a0"], 0)
              np.testing.assert_allclose(res_t["d"]["a1"], i)
              np.testing.assert_allclose(res_t["d"]["a2"], 2)
              np.testing.assert_allclose(res_t["t4"][0], 4)
              np.testing.assert_allclose(res_t["t4"][1], i)
  
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 7 for calling the TreeUtilTest obj's underlying member method code with example usage: obj.test_tree_split_leaves() without return types