code: def gen_function(size, in_types):
      eqns = []
      in_vars = map(fresh_var, in_types)
      cur_vars = in_vars[:]
      for _ in range(gen_nonneg_int(size)):
          if not cur_vars:
              break
          if npr.rand() < subfun_prob:
              arg_vars = gen_subset(cur_vars)
              arg_types = [v.vartype for v in arg_vars]
              fun, out_types = gen_function(
                  size / size_reduction_factor, arg_types
              )
              fun = partial(eval_fun, fun)
              fun = maybe_jit(fun, len(arg_types))
          else:
              arity = choice(list(primitive_generators))
              arg_vars = gen_sized_subset(cur_vars, arity)
              arg_types = [v.vartype for v in arg_vars]
              prim_gen = weighted_choice(primitive_generators[arity])
              fun, out_type = prim_gen(size, *arg_types)
              fun = wrap_singleton(fun)
              out_types = [out_type]
  
          out_vars = map(fresh_var, out_types)
          eqns.append(Eqn(arg_vars, out_vars, fun))
          cur_vars.extend(out_vars)
          cur_vars = thin(cur_vars, thin_prob)
  
      out_vars = gen_subset(cur_vars)
      return Fun(in_vars, out_vars, eqns), [v.vartype for v in out_vars]
  
imports: from collections import namedtuple
  from functools import partial
  import numpy.random
  from absl.testing import absltest
  from absl.testing import parameterized
  import itertools
  import jax.numpy
  from jax import jit, jvp, vjp
  import jax._src.test_util
  from jax import config
  from jax._src.util import unzip2, safe_zip, safe_map
label_desc: Use image node_id 2 for calling a global function with example usage: gen_function(size, in_types) and returns: Fun,  unknown