code: def test_tree_split_key(self):
      rng_key = jax.random.PRNGKey(42)
      tree_like = (1, (2, 3), {"a": 4})
      _, tree_keys = tree_util.tree_split_key(rng_key, tree_like)
      self.assertLen(jax.tree_leaves(tree_keys), 4)
  
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 1 for calling the TreeUtilTest obj's underlying member method code with example usage: obj.test_tree_split_key() without return types