Add `experimental_distribute_datasets_from_function` callback

[khatchad]

Blocked on #115.

          For [`experimental_distribute_datasets_from_function`](https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/distribute/Strategy#experimental_distribute_values_from_function), that function:

Generates tf.distribute.DistributedValues from value_fn.

Tracking DistributedValues isn't currently supported.

Originally posted by @khatchad in https://github.com/wala/ML/issues/92#issuecomment-1839575322

[khatchad]

Some work on this here by @tatianacv:

• https://github.com/ponder-lab/ML/pull/57/files#diff-679e457ba33ce5658ad9444fd41a8839e363bac4c0a268b27a86bbaad0c22056R797-R803
• https://github.com/ponder-lab/ML/pull/57/files#diff-679e457ba33ce5658ad9444fd41a8839e363bac4c0a268b27a86bbaad0c22056R811-R833

[khatchad]

Additional references:

• https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/distribute/MirroredStrategy#distribute_datasets_from_function
• https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/distribute/MirroredStrategy#experimental_distribute_values_from_function
• https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/distribute/Strategy#distribute_datasets_from_function
• https://www.tensorflow.org/versions/r2.9/api_docs/python/tf/distribute/Strategy#experimental_distribute_values_from_function

[khatchad]

The following references may help with the test case, and this is most likely also blocked on #89:

• https://www.tensorflow.org/api_docs/python/tf/data/Dataset#from_tensor_slices
• https://www.tensorflow.org/api_docs/python/tf/data/Iterator

[khatchad]

Dupe of #113.