In multiple cases (PLUMED integration, lammps' compute metatomic) it is required to know the shape of the output before we even can call the model. In PLUMED we end up calling the model with 0 atoms just to extract this information, but in general it would be useful to have an API for this on the models, maybe something like
class MyModel:
def output_layout(self, outputs: Dict[str, ModelOutput]) -> Dict[str, TensorMap]:
...
With the function returning 0-samples TensorMap with the correct metadata. We could even provide a function to implement this easily for outputs with statically known shapes (i.e. most standard outputs)
In multiple cases (PLUMED integration, lammps'
compute metatomic) it is required to know the shape of the output before we even can call the model. In PLUMED we end up calling the model with 0 atoms just to extract this information, but in general it would be useful to have an API for this on the models, maybe something likeWith the function returning 0-samples TensorMap with the correct metadata. We could even provide a function to implement this easily for outputs with statically known shapes (i.e. most standard outputs)