code: def _register_pytree_with_jax(
      pytree_type: type,
      flatten_fn: FlattenFn,
      unflatten_fn: UnflattenFn,
  ):
      def jax_unflatten(aux, parameters):
          return unflatten_fn(parameters, aux)
  
      jax_tree_util.register_pytree_node(
          pytree_type, flatten_fn, jax_unflatten
      )
  
label_desc: Use image node_id 1 for calling a global function with example usage: _register_pytree_with_jax(pytree_type, flatten_fn, unflatten_fn) and returns: unflatten_fn