code: def _maybe_weight_observation_and_reward(
      observation, reward, weights
  ):
      if weights is None:
          return (observation, reward)
      else:
          w_sqrt = tf.sqrt(weights)
          return (
              tf.reshape(w_sqrt, [-1, 1]) * observation,
              w_sqrt * reward,
          )
  
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 7 for calling a global function with example usage: _maybe_weight_observation_and_reward(observation, reward, weights) and returns: observation,  reward, unknown,  unknown