Add Edge Device Support (focus iOS)

FedJAX is currently excellent for simulation-based federated learning research on GPUs/TPUs. However, it lacks edge device support, which is increasingly important for real-world deployment of federated learning.

I’d like to propose adding initial support and pathways for federated training on edge devices, with an emphasis on iOS (Metal-accelerated training) as a starting point.