Training example error

[Youhaojen]

Dear Mario and Ilyes,

Thanks for the nice work. The error occurred when I trained the example.

crude-dorena 2024-08-07 13:15:32.649 INFO: MACE version: 0.1.1
crude-dorena 2024-08-07 13:15:32.895 INFO: Loaded 1200 training configurations from 'rmd17_aspirin_train.xyz'
crude-dorena 2024-08-07 13:15:32.896 INFO: Using random 20.8% of training set for validation
crude-dorena 2024-08-07 13:15:33.253 INFO: Loaded 1800 test configurations from 'rmd17_aspirin_test.xyz'
crude-dorena 2024-08-07 13:15:33.254 INFO: Total number of configurations: train=951, valid=249, test=1800
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 951/951 [00:00<00:00, 15684.05it/s]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 249/249 [00:00<00:00, 15888.97it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1800/1800 [00:00<00:00, 15973.67it/s]
crude-dorena 2024-08-07 13:15:33.455 INFO: z_table= AtomicNumberTable: (1, 6, 8)
crude-dorena 2024-08-07 13:15:33.467 INFO: Compute the average number of neighbors: 14.370
crude-dorena 2024-08-07 13:15:33.468 INFO: Do not normalize the radial basis (avg_r_min=None)
crude-dorena 2024-08-07 13:15:33.473 INFO: Computed average Atomic Energies using least squares: {1: 0.0, 6: 0.0, 8: 0.0}
crude-dorena 2024-08-07 13:15:33.473 INFO: Create MACE with parameters {'symmetric_tensor_product_basis': True, 'off_diagonal': False, 'max_ell': 3, 'hidden_irreps': '256x0e + 256x1o', 'interaction_irreps': 'o3_restricted', 'epsilon': 0.4, 'correlation': 3, 'gate': <function silu at 0x751e410169e0>, 'readout_mlp_irreps': '16x0e', 'r_max': 5.0, 'avg_num_neighbors': 14.370236843422964, 'num_interactions': 2, 'avg_r_min': None, 'num_species': 11, 'radial_basis': <function bessel_basis at 0x751de95903a0>, 'radial_envelope': <function soft_envelope at 0x751de95904c0>}
crude-dorena 2024-08-07 13:15:33.678 INFO: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
crude-dorena 2024-08-07 13:15:33.679 INFO: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
2024-08-07 13:15:33.793841: W external/xla/xla/service/gpu/nvptx_compiler.cc:836] The NVIDIA driver's CUDA version is 12.2 which is older than the PTX compiler version (12.6.20). Because the driver is older than the PTX compiler version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.
crude-dorena 2024-08-07 13:15:34.070 INFO: model: num_features=256 hidden_irreps=1x0e+1x1o interaction_irreps=1x0e+1x1o+1x2e+1x3o 
crude-dorena 2024-08-07 13:15:53.262 INFO: Number of parameters: 5287440
crude-dorena 2024-08-07 13:15:53.262 INFO: Number of parameters in optimizer: 15862322
crude-dorena 2024-08-07 13:15:53.263 INFO: Started training
eval_valid:   0%|                                                                                                                          | 0/50 [00:00<?, ?it/s]crude-dorena 2024-08-07 13:16:11.711 INFO: Compiled function `predictor` for args:
crude-dorena 2024-08-07 13:16:11.711 INFO: - n_node=[21 21 21 21 21  7] total=112
crude-dorena 2024-08-07 13:16:11.711 INFO: - n_edge=[304 302 302 300 304  24] total=1536
crude-dorena 2024-08-07 13:16:11.712 INFO: cache size: 1
eval_valid:   0%|                                                                                                                          | 0/50 [00:19<?, ?it/s]
Traceback (most recent call last):
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/run_train.py", line 76, in <module>
    main()
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/run_train.py", line 58, in main
    train(
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/gin/config.py", line 1605, in gin_wrapper
    utils.augment_exception_message_and_reraise(e, err_str)
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/gin/utils.py", line 41, in augment_exception_message_and_reraise
    raise proxy.with_traceback(exception.__traceback__) from None
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/gin/config.py", line 1582, in gin_wrapper
    return fn(*new_args, **new_kwargs)
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/tools/gin_functions.py", line 312, in train
    loss_ = eval_and_print(valid_loader, "eval_valid")
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/tools/gin_functions.py", line 239, in eval_and_print
    loss_, metrics_ = tools.evaluate(
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/tools/train.py", line 179, in evaluate
    loss_fn(
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/modules/loss.py", line 51, in __call__
    loss += self.forces_weight * mean_squared_error_forces(graph, forces)
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/modules/loss.py", line 22, in mean_squared_error_forces
    graph, jnp.mean(jnp.square(forces_ref - forces_pred), axis=1)
TypeError: unsupported operand type(s) for -: 'NoneType' and 'jaxlib.xla_extension.ArrayImpl'
  In call to configurable 'train' (<function train at 0x751de95901f0>)

Could you please help me fix the error?

By the way, the environment was built by following these steps. Maybe this information will be useful for someone.

pip install -U "jax[cuda12]"
pip install optax
pip install numpy==1.26.4
pip install dm-haiku
pip install git+https://github.com/ACEsuit/mace-jax

Define the NumPy version as 1.26.4 because of the following information.

A module that was compiled using NumPy 1.x cannot be run in
NumPy 2.0.1 as it may crash. To support both 1.x and 2.x
versions of NumPy, modules must be compiled with NumPy 2.0.
Some module may need to rebuild instead e.g. with 'pybind11>=2.12'.

If you are a user of the module, the easiest solution will be to
downgrade to 'numpy<2' or try to upgrade the affected module.
We expect that some modules will need time to support NumPy 2.

Traceback (most recent call last):  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/run_train.py", line 8, in <module>
    from mace_jax import tools
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/tools/__init__.py", line 26, in <module>
    from .train import evaluate, train
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/tools/train.py", line 14, in <module>
    from mace_jax import data, tools
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/data/__init__.py", line 1, in <module>
    from .neighborhood import get_neighborhood
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/data/neighborhood.py", line 3, in <module>
    import matscipy.neighbours
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/matscipy/neighbours.py", line 38, in <module>
    from . import ffi
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/matscipy/ffi.py", line 36, in <module>
    from ._matscipy import *  # noqa
AttributeError: _ARRAY_API not found
Traceback (most recent call last):
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/run_train.py", line 8, in <module>
    from mace_jax import tools
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/tools/__init__.py", line 26, in <module>
    from .train import evaluate, train
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/tools/train.py", line 14, in <module>
    from mace_jax import data, tools
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/data/__init__.py", line 1, in <module>
    from .neighborhood import get_neighborhood
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/mace_jax/data/neighborhood.py", line 3, in <module>
    import matscipy.neighbours
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/matscipy/neighbours.py", line 38, in <module>
    from . import ffi
  File "/home/danken/application/compiler/anaconda3/envs/MACE_JAX/lib/python3.10/site-packages/matscipy/ffi.py", line 36, in <module>
    from ._matscipy import *  # noqa
ImportError: numpy.core.multiarray failed to import

Best regards,
Hao-Jen You

[Clivia117]

Hi. I met same problems as you, it caused by problem of read dataset, the forces_ref is None in original script.
you can try change model of config_from_atoms in mace-jax/data/util.py with changing energy = atoms.info.get(energy_key, None) to energy = atoms.get_potential_energy() and forces = atoms.arrays.get(forces_key, None) to forces = atoms.get_forces().

best regards.
lijun

[ilyes319]

Alternatively, you can downgrade your ASE version to an older one. This is due a recent change in ASE.