code: def _get_initial_and_final_steps_with_action_mask(
      batch_size, context_dim, num_actions
  ):
      observation = np.array(range(batch_size * context_dim)).reshape(
          [batch_size, context_dim]
      )
      observation = tf.constant(observation, dtype=tf.float32)
      mask = 1 - tf.eye(
          batch_size, num_columns=num_actions, dtype=tf.int32
      )
      reward = np.random.uniform(0.0, 1.0, [batch_size])
      initial_step = ts.TimeStep(
          tf.constant(
              ts.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, mask),
      )
      final_step = ts.TimeStep(
          tf.constant(
              ts.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 + 100.0, mask),
      )
      return initial_step, final_step
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow
  from tf_agents.bandits.agents import dropout_thompson_sampling_agent
  from tf_agents.bandits.drivers import driver_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
label_desc: Use image node_id 2 for calling a global function with example usage: _get_initial_and_final_steps_with_action_mask(batch_size, context_dim, num_actions) and returns: initial_step,  final_step