code: def testRestartRewardDTypeArray(self):
      observation = -1
      observation_batch = tf.random.uniform([5, 2])
      reward_spec = array_spec.ArraySpec(shape=[2], dtype=np.float16)
      time_step = ts.restart(observation, None, reward_spec).reward
      time_step_batch = ts.restart(
          observation_batch, 5, reward_spec
      ).reward
  
      self.assertEqual(time_step.dtype, reward_spec.dtype)
      self.assertEqual(time_step_batch.dtype, reward_spec.dtype)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import sys
  import numpy
  import tensorflow
  from tf_agents.specs import array_spec
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 2 for calling the TimeStepTest obj's underlying member method code with example usage: obj.testRestartRewardDTypeArray() without return types