Tests

Add support for `DistributedEmbedding` for Ironwood and expose `has_sparsecores`. #1049

Job	Run time
Test the code on TPU (jax)	6m 20s
Check the code format	1m 58s
Test the code on GPU (torch)	4m 22s
Test the code on TPU (tensorflow)	4m 31s
Test the code on GPU (jax)	5m 54s
Test the code on CPU (tensorflow)	2m 59s
Test the code on CPU (torch)	3m 2s
Test the code on GPU (tensorflow)	5m 47s
Test the code on CPU (jax)	7m 19s
	42m 12s

Provide feedback