code: def test_psum_unary_function_applies_psum_to_outputs(self):
      def f(x):
          return x
  
      f = distribute_lib.make_psum_function(
          f, self.axis_name, self.axis_name, out_dtype=tf.float32
      )
  
      x = self.shard_values(tf.ones(4))
      out_parts = self.per_replica_to_tensor(self.strategy_run(f, (x,)))
  
      self.assertAllEqual(
          self.evaluate(out_parts), self.evaluate(4 * tf.ones(4))
      )
  
imports: import functools
  import itertools
  from absl.testing import parameterized
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 1 for calling the ShardedFunctionTest obj's underlying member method code with example usage: obj.test_psum_unary_function_applies_psum_to_outputs() and returns: x