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()

@@ -1,12 +1,375 @@
+# Imports the Cloud Logging client library%0Aimport google.cloud.logging%0A%0A# Instantiates a client%0Aclient = google.cloud.logging.Client()%0A%0A# Retrieves a Cloud Logging handler based on the environment%0A# you're running in and integrates the handler with the%0A# Python logging module. By default this captures all logs%0A# at INFO level and higher%0Aclient.setup_logging()%0A%0A
 import tenso
@@ -424,16 +424,17 @@
 dtensor%0A
+%0A
 print('c

enable cloud logging.