code: def test_psum_binary_function_applies_psum_to_outputs(self):
      def f(x, y):
          return x + y
  
      f_psum = distribute_lib.make_psum_function(
          f,
          (self.axis_name, self.axis_name),
          self.axis_name,
          out_dtype=tf.float32,
      )
  
      x = self.shard_values(tf.ones(4))
      y = self.shard_values(2 * tf.ones(4))
      out_parts = self.per_replica_to_tensor(
          self.strategy_run(f_psum, (x, y))
      )
  
      self.assertAllEqual(
          self.evaluate(out_parts), self.evaluate(12 * tf.ones(4))
      )
  
      f_psum = distribute_lib.make_psum_function(
          f,
          (self.axis_name, self.axis_name),
          None,
          out_dtype=tf.float32,
      )
  
      x = self.shard_values(tf.ones(4))
      y = self.shard_values(2 * tf.ones(4))
      out_parts = self.per_replica_to_tensor(
          self.strategy_run(f_psum, (x, y))
      )
  
      self.assertAllEqual(
          self.evaluate(out_parts), self.evaluate(3 * tf.ones(4))
      )
  
      f_psum = distribute_lib.make_psum_function(
          f,
          (self.axis_name, self.axis_name),
          None,
          out_dtype=tf.float32,
      )
  
      x = self.shard_values(tf.ones(4))
      y = self.shard_values(2 * tf.ones(4))
      out_parts = self.per_replica_to_tensor(
          self.strategy_run(f_psum, (x, y), in_axes=(0, 0))
      )
  
      self.assertAllEqual(
          self.evaluate(out_parts), self.evaluate(3 * tf.ones(4))
      )
  
      f_psum = distribute_lib.make_psum_function(
          f, (self.axis_name, None), None, out_dtype=tf.float32
      )
  
      x = self.shard_values(tf.ones(4))
      y = 2.0
      out_parts = self.per_replica_to_tensor(
          self.strategy_run(f_psum, (x, y), in_axes=(0, None))
      )
  
      self.assertAllEqual(
          self.evaluate(out_parts), self.evaluate(3 * 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 2 for calling the ShardedFunctionTest obj's underlying member method code with example usage: obj.test_psum_binary_function_applies_psum_to_outputs() and returns: unknown