wala / ML

Eclipse Public License 2.0
23 stars 17 forks source link

Add Strategy.reduce() to XML summaries #111

Closed khatchad closed 7 months ago

khatchad commented 7 months ago

The API at https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/distribute/Strategy#reduce returns a tensor. Can we track it? We need a summary for the return value (it's a Tensor):

From the docs:

Returns A Tensor that's the concatenation of value across replicas along axis dimension.

Example from the docs:

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
# A DistributedValues with component tensor of shape (2, 1) on each replica
distributed_values = strategy.experimental_distribute_values_from_function(lambda _: tf.identity(tf.constant([[1], [2]])))
@tf.function
def run():
  return strategy.gather(distributed_values, axis=0)
run()
khatchad commented 7 months ago

I don't think this is a "generator."