.. _multihost:

.. module:: mlip.utils.multihost

Multi-host utilities
====================

.. autofunction:: create_device_mesh

.. autofunction:: create_replicated_sharding

.. autofunction:: create_dp_sharding

.. autofunction:: sync_string

.. autofunction:: only_specific_processes

.. autofunction:: single_host_jax_and_orbax

.. autofunction:: assert_pytrees_match_across_hosts

.. autofunction:: assert_pytrees_match