code: def testTimeMajorBatchMajorDiscountedReturnsAreSame(
      self, num_time_steps, batch_size, with_final_value
  ):
      rewards = np.random.rand(num_time_steps, batch_size).astype(
          np.float32
      )
      discounts = np.random.rand(num_time_steps, batch_size).astype(
          np.float32
      )
      final_value = (
          np.random.rand(batch_size).astype(np.float32)
          if with_final_value
          else None
      )
  
      time_major_discounted_return = value_ops.discounted_return(
          rewards=rewards, discounts=discounts, final_value=final_value
      )
  
      batch_major_discounted_return = value_ops.discounted_return(
          rewards=tf.transpose(rewards),
          discounts=tf.transpose(discounts),
          final_value=final_value,
          time_major=False,
      )
  
      self.assertAllClose(
          time_major_discounted_return,
          tf.transpose(batch_major_discounted_return),
      )
  
      single_time_major_discounted_return = value_ops.discounted_return(
          rewards=rewards,
          discounts=discounts,
          final_value=final_value,
          provide_all_returns=False,
      )
  
      single_batch_major_discounted_return = (
          value_ops.discounted_return(
              rewards=tf.transpose(rewards),
              discounts=tf.transpose(discounts),
              final_value=final_value,
              time_major=False,
              provide_all_returns=False,
          )
      )
  
      self.assertAllClose(
          single_time_major_discounted_return,
          time_major_discounted_return[0],
      )
      self.assertAllClose(
          single_batch_major_discounted_return,
          time_major_discounted_return[0],
      )
  
file_path: agents/tf_agents/utils/value_ops_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents.utils import value_ops
label_desc: Use image node_id 2 for calling the DiscountedReturnTest obj's underlying member method code with example usage: obj.testTimeMajorBatchMajorDiscountedReturnsAreSame(num_time_steps, batch_size, with_final_value) without return types