float128 error at import

Hi,

I’ve got a mac os M1 running, and at import mpi4jax I’m getting the following. (fyi, running mpi4jax 0.3.5)

File /opt/homebrew/Caskroom/miniforge/base/envs/localjw/lib/python3.9/site-packages/mpi4jax/__init__.py:9, in <module>
      6 __version__ = get_versions()["version"]
      7 del get_versions
----> 9 from ._src import (  # noqa: E402
     10     allgather,
     11     allreduce,
     12     alltoall,
     13     barrier,
     14     bcast,
     15     gather,
     16     recv,
     17     reduce,
     18     scan,
     19     scatter,
     20     send,
     21     sendrecv,
     22     has_cuda_support,
     23 )
     25 __all__ = [
     26     "allgather",
     27     "allreduce",
   (...)
     38     "has_cuda_support",
     39 ]

File /opt/homebrew/Caskroom/miniforge/base/envs/localjw/lib/python3.9/site-packages/mpi4jax/_src/__init__.py:9, in <module>
      6 from . import xla_bridge
      8 # import public API
----> 9 from .collective_ops.allgather import allgather  # noqa: F401
     10 from .collective_ops.allreduce import allreduce  # noqa: F401
     11 from .collective_ops.alltoall import alltoall  # noqa: F401

File /opt/homebrew/Caskroom/miniforge/base/envs/localjw/lib/python3.9/site-packages/mpi4jax/_src/collective_ops/allgather.py:10, in <module>
      7 from jax.lax import create_token
      8 from jax.lib import xla_client
---> 10 from ..utils import (
     11     HashableMPIType,
     12     default_primitive_impl,
     13     to_dtype_handle,
     14     to_mpi_handle,
     15     unpack_hashable,
     16     wrap_as_hashable,
     17     xla_constant_intc,
     18     xla_constant_uintptr,
     19 )
     20 from ..decorators import translation_rule_cpu, translation_rule_gpu
     21 from ..validation import enforce_types

File /opt/homebrew/Caskroom/miniforge/base/envs/localjw/lib/python3.9/site-packages/mpi4jax/_src/utils.py:45, in <module>
     36     """
     37     Returns a pointer to the underlying C MPI object
     38     """
     39     return _np.uintp(_MPI._addressof(mpi_obj))
     42 MPI_TYPE_MAP = {
     43     _np.dtype(_np.float32): _MPI.FLOAT,
     44     _np.dtype(_np.float64): _MPI.DOUBLE,
---> 45     _np.dtype(_np.float128): _MPI.LONG_DOUBLE,
     46     _np.dtype(_np.complex64): _MPI.COMPLEX,
     47     _np.dtype(_np.complex128): _MPI.DOUBLE_COMPLEX,
     48     _np.dtype(_np.int8): _MPI.INT8_T,
     49     _np.dtype(_np.int16): _MPI.INT16_T,
     50     _np.dtype(_np.int32): _MPI.INT32_T,
     51     _np.dtype(_np.int64): _MPI.INT64_T,
     52     _np.dtype(_np.uint8): _MPI.UINT8_T,
     53     _np.dtype(_np.uint16): _MPI.UINT16_T,
     54     _np.dtype(_np.uint32): _MPI.UINT32_T,
     55     _np.dtype(_np.uint64): _MPI.UINT64_T,
     56     _np.dtype(_np.bool_): _MPI.BOOL,
     57 }
     60 def to_dtype_handle(dtype):
     61     """
     62     Returns the pointer to the MPI dtype of the input numpy dtype
     63     """

File /opt/homebrew/Caskroom/miniforge/base/envs/localjw/lib/python3.9/site-packages/numpy/__init__.py:313, in __getattr__(attr)
    310     from .testing import Tester
    311     return Tester
--> 313 raise AttributeError("module {!r} has no attribute "
    314                      "{!r}".format(__name__, attr))

AttributeError: module 'numpy' has no attribute 'float128'

[jwnys]

Perfect, works now. Thanks for the quick fix!

[Thenerdstation]

I’m swamped with onsite interviews for the next 2 weeks so I wouldn’t expect to have a mergeable agreed upon solution before then. Do what you will given that.