Skip to content

furax.distributed

Functions:

furax.distributed.maybe_init()

Source code in src/furax/distributed.py 
 5
 6
 7
 8
 9
10
11
def maybe_init() -> bool:
    try:
        jax.distributed.initialize()
    except ValueError:
        return False
    mhu.sync_global_devices('furax.startup')
    return True