import tensorflow as tf
from tensorflow.experimental import dtensor
print('client', dtensor.client_id(), 'devices', tf.config.list_physical_devices('GPU'))
print(tf.__version__)
dtensor.initialize_multi_client()
mesh = dtensor.create_distributed_mesh([("batch", 4)], device_type='GPU', num_global_devices=4)

def main():
    layout = dtensor.Layout(['batch', dtensor.UNSHARDED], mesh)
    data = dtensor.call_with_layout(tf.ones, layout, shape=(4, 100))
    print(data)
    print(tf.reduce_sum(data))

main()

# Imports the Cloud Logging client library
import google.cloud.logging

# Instantiates a client
client = google.cloud.logging.Client()

# Retrieves a Cloud Logging handler based on the environment
# you're running in and integrates the handler with the
# Python logging module. By default this captures all logs
# at INFO level and higher
client.setup_logging()

import tensorflow as tf
from tensorflow.experimental import dtensor

print('client', dtensor.client_id(), 'devices', tf.config.list_physical_devices('GPU'))
print(tf.__version__)
dtensor.initialize_multi_client()
mesh = dtensor.create_distributed_mesh([("batch", 4)], device_type='GPU', num_global_devices=4)

def main():
    layout = dtensor.Layout(['batch', dtensor.UNSHARDED], mesh)
    data = dtensor.call_with_layout(tf.ones, layout, shape=(4, 100))
    print(data)
    print(tf.reduce_sum(data))

main()

enable cloud logging.