Is support for MACE-POLAR implemented in eval_configs.py? If so, how to activate it?

[godot11]

I'm trying to run eval_configs.py on a dataset I have previously managed to process with foundation models as

python3 "${MACE_DIR}/mace/cli/eval_configs.py"
--configs="configs.xyz"
--model="MACE-matpes-pbe-omat-ft.model"
--output="configs_processed.xyz"
--device="cuda"
--batch_size=6

However, with the polar MACE models, such as MACE-POLAR-1-M.model, I get the following error:

Traceback (most recent call last):
  File "[HOME]/Projects/MACE_training/mace/mace/cli/eval_configs.py", line 337, in <module>
    main()
  File "[HOME]/Projects/MACE_training/mace/mace/cli/eval_configs.py", line 132, in main
    run(args)
  File "[HOME]/Projects/MACE_training/mace/mace/cli/eval_configs.py", line 195, in run
    output = get_model_output(
             ^^^^^^^^^^^^^^^^^
  File "[HOME]/Projects/MACE_training/mace/mace/cli/eval_configs.py", line 127, in get_model_output
    return model(batch, **forward_args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[HOME]/micromamba/envs/torch311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[HOME]/micromamba/envs/torch311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[HOME]/micromamba/envs/torch311/lib/python3.11/site-packages/mace/modules/extensions.py", line 650, in forward
    node_feats = self.node_embedding(data["node_attrs"])
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[HOME]/micromamba/envs/torch311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[HOME]/micromamba/envs/torch311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[HOME]/micromamba/envs/torch311/lib/python3.11/site-packages/mace/modules/blocks.py", line 57, in forward
    return self.linear(node_attrs)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "[HOME]/micromamba/envs/torch311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[HOME]/micromamba/envs/torch311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[HOME]/micromamba/envs/torch311/lib/python3.11/site-packages/e3nn/o3/_linear.py", line 280, in forward
    return self._compiled_main(features, weight, bias)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[HOME]/micromamba/envs/torch311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[HOME]/micromamba/envs/torch311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/torch/fx/graph_module/___torch_mangle_55.py", line 26, in forward
    reshape_3 = torch.reshape(reshape_1, [83, 512])
    mul = torch.mul(reshape_2, 0.10976425998969035)
    tensordot = _0(mul, reshape_3, ([1], [0]), None, )
                ~~ <--- HERE
    permute = torch.permute(tensordot, [0, 2, 1])
    reshape_4 = torch.reshape(permute, [getitem_1, 512])
  File "code/__torch__/torch/functional/___torch_mangle_0.py", line 9, in tensordot
  dims_a, dims_b, = dims1
  if torch.__is__(out, None):
    _1 = torch.tensordot(a, b, dims_a, dims_b)
         ~~~~~~~~~~~~~~~ <--- HERE
    _0 = _1
  else:

Traceback of TorchScript, original code (most recent call last):
  File "<eval_with_key>.404", line 15, in forward
    reshape_3 = reshape_1.reshape((83, 512));  reshape_1 = None
    mul = reshape_2 * 0.10976425998969035;  reshape_2 = None
    tensordot = torch.functional.tensordot(mul, reshape_3, ([1], [0]), out = None);  mul = reshape_3 = None
                ~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    permute = tensordot.permute([0, 2, 1]);  tensordot = None
    reshape_4 = permute.reshape(getitem_1, 512);  permute = getitem_1 = None
  File "/Users/Lilyes/miniconda3/envs/torch_env/lib/python3.11/site-packages/torch/functional.py", line 1382, in tensordot

    if out is None:
        return _VF.tensordot(a, b, dims_a, dims_b)  # type: ignore[attr-defined]
               ~~~~~~~~~~~~~ <--- HERE
    else:
        return _VF.tensordot(a, b, dims_a, dims_b, out=out)  # type: ignore[attr-defined]
RuntimeError: both inputs should have same dtype

Is support for PolarMACE not implemented in eval_configs.py, or did I miss a flag?
If it would be helpful, I will submit a proper bug report.

[ilyes319]

Hey @godot11, can you try specifying "float32" as input dtype to the eval script? it looks like that's the problem.

[godot11]

Hey @ilyes319, thanks - using --default_dtype=float32 indeed resolved the issue. Using --default_dtype=float64 reproduces the error.

I'm a bit confused though. I have realized that while other foundation models load as float64 models (model = torch.load(f"model_name.model"); print(next(model.parameters()).dtype)), all polar models load as float32. This likely caused my issue. However, finetuning a polar foundation model with --default_dtype=float64 is suggested by the documentation, and such a finetuning indeed produces a seemingly proper float64 model.

Are polar models float32 models, or float64 models that appear as float32 due to some technicality? Is it okay to finetune them with --default_dtype=float64 and obtain a float64 model, effectively "upscaling" their output resolution?

[Edit] I raised this in ACEsuit/mace-foundations#69 approprietly.