code: def _get_initial_and_final_steps_with_per_arm_features(
      batch_size,
      global_context_dim,
      num_actions,
      arm_context_dim,
      apply_action_mask=False,
      num_actions_feature=False,
  ):
      global_observation = np.array(
          range(batch_size * global_context_dim)
      ).reshape([batch_size, global_context_dim])
      arm_observation = np.array(
          range(batch_size * num_actions * arm_context_dim)
      ).reshape([batch_size, num_actions, arm_context_dim])
      reward = np.random.uniform(0.0, 1.0, [batch_size])
      observation = {
          "global": tf.constant(
              global_observation,
              dtype=tf.float32,
              shape=[batch_size, global_context_dim],
              name="global_observation",
          ),
          "per_arm": tf.constant(
              arm_observation,
              dtype=tf.float32,
              shape=[batch_size, num_actions, arm_context_dim],
              name="arm_observation",
          ),
      }
      if num_actions_feature:
          observation.update(
              {
                  "num_actions": tf.ones([batch_size], dtype=tf.int32)
                  * (num_actions - 1)
              }
          )
      if apply_action_mask:
          observation = (
              observation,
              tf.ones([batch_size, num_actions], dtype=tf.int32),
          )
      initial_step = time_step.TimeStep(
          tf.constant(
              time_step.StepType.FIRST,
              dtype=tf.int32,
              shape=[batch_size],
              name="step_type",
          ),
          tf.constant(
              0.0, dtype=tf.float32, shape=[batch_size], name="reward"
          ),
          tf.constant(
              1.0, dtype=tf.float32, shape=[batch_size], name="discount"
          ),
          observation,
      )
      observation = {
          "global": tf.constant(
              global_observation + 100.0,
              dtype=tf.float32,
              shape=[batch_size, global_context_dim],
              name="global_observation",
          ),
          "arm": tf.constant(
              arm_observation + 100.0,
              dtype=tf.float32,
              shape=[batch_size, num_actions, arm_context_dim],
              name="arm_observation",
          ),
      }
      if num_actions_feature:
          observation.update(
              {
                  "num_actions": tf.ones([batch_size], dtype=tf.int32)
                  * num_actions
              }
          )
      if apply_action_mask:
          observation = (
              observation,
              tf.ones([batch_size, num_actions], dtype=tf.int32),
          )
      final_step = time_step.TimeStep(
          tf.constant(
              time_step.StepType.LAST,
              dtype=tf.int32,
              shape=[batch_size],
              name="step_type",
          ),
          tf.constant(
              reward,
              dtype=tf.float32,
              shape=[batch_size],
              name="reward",
          ),
          tf.constant(
              1.0, dtype=tf.float32, shape=[batch_size], name="discount"
          ),
          observation,
      )
      return initial_step, final_step
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from typing import Optional, Tuple
  from absl import flags
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.agents import lin_ucb_agent
  from tf_agents.bandits.agents import linear_bandit_agent
  from tf_agents.bandits.agents import linear_thompson_sampling_agent
  from tf_agents.bandits.agents import utils
  from tf_agents.bandits.drivers import driver_utils
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 3 for calling a global function with example usage: _get_initial_and_final_steps_with_per_arm_features(batch_size, global_context_dim, num_actions, arm_context_dim, apply_action_mask, num_actions_feature) and returns: initial_step,  final_step