code: def get_reward_layer(num_actions=5, encoding_dim=10):
      return tf.keras.layers.Dense(
          num_actions,
          activation=None,
          kernel_initializer=tf.constant_initializer(
              np.ones([encoding_dim, num_actions])
          ),
          bias_initializer=tf.constant_initializer(
              np.array(range(num_actions))
          ),
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.bandits.networks import global_and_arm_feature_network
  from tf_agents.bandits.policies import neural_linucb_policy
  from tf_agents.bandits.specs import utils
  from tf_agents.networks import network
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import common
  from tf_agents.utils import test_utils
label_desc: Use image node_id 1 for calling a global function with example usage: get_reward_layer(num_actions, encoding_dim) and returns: tf