code: def test_tree_fn_with_partials(self):
      def add(x, y):
          return x + y
  
      tree_add = tree_util.tree_fn(add, y=1)
      tree_x = [1, 2, 3]
      output = tree_add(tree_x)
      exp_output = [2, 3, 4]
      # Test output.
      np.testing.assert_allclose(output, exp_output)
  
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 9 for calling the TreeUtilTest obj's underlying member method code with example usage: obj.test_tree_fn_with_partials() and returns: unknown