code: def test_can_shard_values_across_logical_devices(self):
      @tf.function(autograph=False)
      def value_fn(ctx):
          return tf.cast(ctx.replica_id_in_sync_group, tf.float32)
  
      def add_one(x):
          return x + 1.0
  
      values = (
          self.strategy().experimental_distribute_values_from_function(
              value_fn
          )
      )
      out_values = self.evaluate(
          self.per_replica_to_tensor(
              self.strategy_run(add_one, (values,))
          )
      )
      self.assertAllEqual(out_values, [1.0, 2.0, 3.0, 4.0])
  
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 LogProbPartsTest obj's underlying member method code with example usage: obj.test_can_shard_values_across_logical_devices() and returns: tf, unknown