code: def random_linear_multiple_reward_fn_generator(
      context_dim, num_actions, num_rewards, squeeze_dims=True
  ):
      """A function that returns `num_actions` linear functions.
  
      For each action, the corresponding linear function has underlying parameters
      of shape [`context_dim`, 'num_rewards']. Optionally, squeeze can be applied.
  
      Args:
        context_dim: Number of parameters per function.
        num_actions: Number of functions returned.
        num_rewards: Numer of rewards we want to generate.
        squeeze_dims: whether to apply squeeze or not.
  
      Returns:
        A list of linear functions.
      """
  
      def _gen_multiple_rewards_for_action():
          params = np.random.rand(context_dim, num_rewards)
          if squeeze_dims:
              return np.squeeze(params)
          else:
              return params
  
      return linear_multiple_reward_fn_generator(
          [
              _gen_multiple_rewards_for_action()
              for _ in range(num_actions)
          ]
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import gin
  import numpy
  from six.moves import range
  import tensorflow
  from tf_agents.bandits.agents import utils
  from tf_agents.bandits.environments import wheel_py_environment
label_desc: Use image node_id 8 for calling a global function with example usage: random_linear_multiple_reward_fn_generator(context_dim, num_actions, num_rewards, squeeze_dims) and returns: linear_multiple_reward_fn_generator, np, params