:modulename:`compose.mace_jax`
============================================

.. currentmodule:: chemtrain.compose.mace_jax

.. automodule:: chemtrain.compose.mace_jax

.. autofunction:: mace_jax_neighborlist

.. autoclass:: JaxMACE
    :members:


Species Mappings
------------------

.. autoclass:: SpeciesMapping
    :members:

.. autoclass:: AtomicNumberMapping
    :members:


Load Foundational Models
-------------------------

.. autofunction:: load_foundational_model