Fix jax mean for 0 counts.

Author: PatReis

README.md

@@ -295,7 +295,6 @@ You can find a [table](training/results/README.md) of common benchmark datasets
 
 Some known issues to be aware of, if using and making new models or layers with `kgcnn`.
 * Jagged or nested Tensors loading into models for PyTorch backend is not working.
-* ForceModel does not support all backends.
 * BatchNormalization layer dos not support padding yet.
 * Keras AUC metric does not seem to work for torch cuda.
 

changelog.md

@@ -16,6 +16,7 @@ causing clashes with built-in functions. We catch defaults to be at least as bac
 * Added simple ragged support for ``train_force.py``
 * Implemented random equivariant initialize for PAiNN
 * Implemented charge and dipole output for HDNNP2nd
+* Implemented jax backend for force models.
 
 
 v4.0.0

kgcnn/backend/_jax.py

@@ -53,7 +53,8 @@ def scatter_reduce_mean(indices, values, shape):
     zeros = jnp.zeros(shape, values.dtype)
     counts = jnp.zeros(shape, values.dtype)
     counts = counts.at[indices].add(jnp.ones_like(values))
-    return zeros.at[indices].add(values)/counts
+    inverse_counts = jnp.nan_to_num(jnp.reciprocal(counts))
+    return zeros.at[indices].add(values)*inverse_counts
 
 
 def scatter_reduce_softmax(indices, values, shape, normalize: bool = False):