def synchronous_embedding_fn(view: datasets.View) -> chex.Array:
    return jax.device_get(embedding_fn(view))
def synchronous_embedding_fn(view: datasets.View) -> chex.Array:
    return jax.device_get(embedding_fn(view))