code: def _create_simple_agent_and_data(
      set_example_weights: bool,
  ) -> Tuple[
      linear_agent.LinearBanditAgent,
      types.NestedTensor,
      Optional[tf.Tensor],
  ]:
      num_actions = 1
      context_dim = 1
      batch_size = 10
      dtype = tf.float32
      # The observation consists of a single constant feature.
      initial_step, final_step = _get_initial_and_final_steps(
          batch_size, context_dim, use_constant_observations=True
      )
      action = [0] * batch_size
      action_step = _get_action_step(action)
      experience = _get_experience(
          initial_step, action_step, final_step
      )
  
      # Construct an agent and perform the update.
      observation_spec = tensor_spec.TensorSpec(
          [context_dim], tf.float32
      )
      time_step_spec = time_step.time_step_spec(observation_spec)
      action_spec = tensor_spec.BoundedTensorSpec(
          dtype=tf.int32, shape=(), minimum=0, maximum=num_actions - 1
      )
      variable_collection = linear_agent.LinearBanditVariableCollection(
          context_dim, num_actions, use_eigendecomp=False, dtype=dtype
      )
  
      agent = linear_agent.LinearBanditAgent(
          exploration_policy=linear_agent.ExplorationPolicy.linear_ucb_policy,
          time_step_spec=time_step_spec,
          action_spec=action_spec,
          variable_collection=variable_collection,
          use_eigendecomp=False,
          tikhonov_weight=0.0,
          dtype=dtype,
      )
  
      weights = (
          tf.linspace(start=1.5, stop=10.5, num=batch_size)
          if set_example_weights
          else None
      )
  
      return agent, experience, weights
  
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 8 for calling a global function with example usage: _create_simple_agent_and_data(set_example_weights) and returns: agent,  experience,  weights