m3gnet for 1 or 2 atoms ValueError with tensorflow==2.10.0

[mhellstr]

If I try to run a prediction on an H2 molecule or an H atom with the M3GNetCalculator I get the below error. Is it possible to run m3gnet for such systems? For example to get a smooth dissociation curve for H2.

from m3gnet.models import M3GNetCalculator, Potential, M3GNet
from ase.io import read
from ase import Atoms

def main():
    potential = Potential(M3GNet.load())
    calc = M3GNetCalculator(potential=potential)
    print("H2")
    atoms = Atoms('H2', positions=[(0, 0, -0.35), (0, 0, 0.35)])
    atoms.set_calculator(calc)
    try:
        atoms.get_potential_energy()
    except Exception as e:
        print(e)

    print("H")
    atoms = Atoms('H', positions=[(0, 0, -0.35)])
    atoms.set_calculator(calc)
    try:
        atoms.get_potential_energy()
    except Exception as e:
        print(e)

if __name__ == '__main__':
    main()

    File "/home/hellstrom/.scm/python/AMS2022.2.venv/lib/python3.8/site-packages/m3gnet/models/_base.py", line 186, in get_efs_tensor  *
        energies = self.get_energies(graph)
    File "/home/hellstrom/.scm/python/AMS2022.2.venv/lib/python3.8/site-packages/m3gnet/models/_base.py", line 261, in get_energies  *
        return self.model(graph)
    File "/home/hellstrom/.scm/python/AMS2022.2.venv/lib/python3.8/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler  **
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_filex91gvqx1.py", line 13, in tf__call
        three_basis = ag__.converted_call(ag__.ld(self).basis_expansion, (ag__.ld(graph),), None, fscope)
    File "/tmp/__autograph_generated_filekf031m2n.py", line 15, in tf__call
        retval_ = ag__.converted_call(ag__.ld(combine_sbf_shf), (ag__.ld(sbf), ag__.ld(shf)), dict(max_n=ag__.ld(self).max_n, max_l=ag__.ld(self).max_l, use_phi=ag__.ld(self).use_phi), fscope)
    File "/tmp/__autograph_generated_filemd4pex2a.py", line 81, in tf__combine_sbf_shf
        ag__.if_stmt((ag__.converted_call(ag__.ld(tf).shape, (ag__.ld(sbf),), None, fscope)[0] == 0), if_body_2, else_body_2, get_state_2, set_state_2, ('do_return', 'retval_'), 2)
    File "/tmp/__autograph_generated_filemd4pex2a.py", line 50, in else_body_2
        expanded_sbf = ag__.converted_call(ag__.ld(tf).repeat, (ag__.ld(sbf),), dict(repeats=ag__.ld(repeats_sbf), axis=1), fscope)

    ValueError: Exception encountered when calling layer "m3g_net" "                 f"(type M3GNet).
    
    in user code:
    
        File "/home/hellstrom/.scm/python/AMS2022.2.venv/lib/python3.8/site-packages/m3gnet/models/_m3gnet.py", line 253, in call  *
            three_basis = self.basis_expansion(graph)
        File "/home/hellstrom/.scm/python/AMS2022.2.venv/lib/python3.8/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler  **
            raise e.with_traceback(filtered_tb) from None
        File "/tmp/__autograph_generated_filekf031m2n.py", line 15, in tf__call
            retval_ = ag__.converted_call(ag__.ld(combine_sbf_shf), (ag__.ld(sbf), ag__.ld(shf)), dict(max_n=ag__.ld(self).max_n, max_l=ag__.ld(self).max_l, use_phi=ag__.ld(self).use_phi), fscope)
        File "/tmp/__autograph_generated_filemd4pex2a.py", line 81, in tf__combine_sbf_shf
            ag__.if_stmt((ag__.converted_call(ag__.ld(tf).shape, (ag__.ld(sbf),), None, fscope)[0] == 0), if_body_2, else_body_2, get_state_2, set_state_2, ('do_return', 'retval_'), 2)
        File "/tmp/__autograph_generated_filemd4pex2a.py", line 50, in else_body_2
            expanded_sbf = ag__.converted_call(ag__.ld(tf).repeat, (ag__.ld(sbf),), dict(repeats=ag__.ld(repeats_sbf), axis=1), fscope)
    
        ValueError: Exception encountered when calling layer "spherical_bessel_with_harmonics" "                 f"(type SphericalBesselWithHarmonics).
        
        in user code:
        
            File "/home/hellstrom/.scm/python/AMS2022.2.venv/lib/python3.8/site-packages/m3gnet/layers/_three_body.py", line 57, in call  *
                return combine_sbf_shf(sbf, shf, max_n=self.max_n, max_l=self.max_l, use_phi=self.use_phi)
            File "/home/hellstrom/.scm/python/AMS2022.2.venv/lib/python3.8/site-packages/m3gnet/utils/_math.py", line 300, in combine_sbf_shf  *
                expanded_sbf = tf.repeat(sbf, repeats=repeats_sbf, axis=1)
        
            ValueError: Dimension 1 in both shapes must be equal, but are 9 and 0. Shapes are [0,9] and [0,0].
        
        
        Call arguments received by layer "spherical_bessel_with_harmonics" "                 f"(type SphericalBesselWithHarmonics):
          • graph=['tf.Tensor(shape=(2, 1), dtype=int32)', 'tf.Tensor(shape=(2,), dtype=float32)', 'None', 'tf.Tensor(shape=(2, 3), dtype=float32)', 'tf.Tensor(shape=(2, 2), dtype=int32)', 'tf.Tensor(shape=(2, 3), dtype=int32)', 'tf.Tensor(shape=(1,), dtype=int32)', 'tf.Tensor(shape=(1,), dtype=int32)', 'tf.Tensor(shape=(2,), dtype=float32)', 'tf.Tensor(shape=(1, 3, 3), dtype=float32)', 'tf.Tensor(shape=(0, 2), dtype=int32)', 'tf.Tensor(shape=(0,), dtype=float32)', 'tf.Tensor(shape=(0,), dtype=float32)', 'tf.Tensor(shape=(0,), dtype=float32)', 'tf.Tensor(shape=(2,), dtype=int32)', 'tf.Tensor(shape=(2,), dtype=int32)', 'tf.Tensor(shape=(1,), dtype=int32)']
          • kwargs={'training': 'None'}
    
    
    Call arguments received by layer "m3g_net" "                 f"(type M3GNet):
      • graph=['tf.Tensor(shape=(2, 1), dtype=int32)', 'tf.Tensor(shape=(2, 1), dtype=float32)', 'None', 'tf.Tensor(shape=(2, 3), dtype=float32)', 'tf.Tensor(shape=(2, 2), dtype=int32)', 'tf.Tensor(shape=(2, 3), dtype=int32)', 'tf.Tensor(shape=(1,), dtype=int32)', 'tf.Tensor(shape=(1,), dtype=int32)', 'tf.Tensor(shape=(2,), dtype=float32)', 'tf.Tensor(shape=(1, 3, 3), dtype=float32)', 'tf.Tensor(shape=(0, 2), dtype=int32)', 'None', 'None', 'None', 'tf.Tensor(shape=(2,), dtype=int32)', 'tf.Tensor(shape=(2,), dtype=int32)', 'tf.Tensor(shape=(1,), dtype=int32)']
      • kwargs={'training': 'None'}

[mhellstr]

I should emphasize that I'm mainly interested in having it work from a technical perspective. I understand that it wasn't part of the training data, so I'm not expecting very accurate predictions.

[chc273]

@mhellstr try tensorflow 2.9.1

[chc273]

Here is what I got for your reference

[mhellstr]

Perfect, thanks, I was using tensorflow 2.10.0 for which the error above appears.