Given that Apple has built PyTorch for Metal, I would assume they handle unified memory properly in there too… Same for Tensorflow and Jax on Metal.