fix runtime errors in colocated python dataloader

[sadikneipp]

Title

Fix: Resolve runtime errors in colocated TFDS dataloader

Description

fixes:

• _tfds_data_processing.py:

  • The tfds.builder now correctly uses the dataset_path variable, which is needed for tfds.
  • The constructor for RemoteIterator has been updated to pass global_shape as argument, fixing an interface mismatch.

• multihost_dataloading.py:

  • The put_to_tpu_devices function now correctly returns the array after it has been placed on the TPU device. Otherwise the iterator would have empty data.
  • Initialization logging for RemoteIterator is now more informative, including a test output to aid in debugging and validation.

Tests

This patch was used in the last round of scale tests, validated up to 4 slice scale.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[aireenmei]

Thanks Sadi!