Distributed Training API

[innat]

1. In this guide, it only talks about jax backend distributed training using keras 3 API. Do other backends also support this? If not, is there any raod-map for that?
2. Using tensorflow.python.client.device_lib, we can get the following information about GPU(s).

from tensorflow.python.client import device_lib

list_ld = device_lib.list_local_devices()
   for dev in list_ld: 
      print(dev.name,dev.memory_limit)
      
/device:CPU:0 268435456
/device:GPU:0 14619377664
/device:GPU:1 14619377664

print('TF System Cuda V.  : ', tf.sysconfig.get_build_info()["cuda_version"])
print('TF System CudNN V. : ', tf.sysconfig.get_build_info()["cudnn_version"])

Would it be convenient to have someting general from keras 3 API, no? In torch, it provides the following, convenient when needed for sanity check.

torch.version.cuda
torch.backends.cudnn.version()
torch.cuda.is_available()
torch.cuda.device_count()   
torch.cuda.current_device()
torch.cuda.device(0)
torch.cuda.get_device_name(0)  
torch.distributed.is_nccl_available()
torch.cuda.nccl.version()