Ensemble do not work with ApaxJaxMD

[PythonFZ]

When using a full ensemble one gets:

ScopeParamShapeError: Initializer expected to generate shape (5, 119, 119, 5, 7) but got shape (119, 119, 5, 7) instead 
for parameter "atomic_type_embedding" in "/energy_model/atomistic_model/descriptor/radial_fn".

reproduce with https://github.com/IPSProjects/gmx-uim/tree/7bcd61d43958e10a2e15ef0785fc5231d92ee87f

When using a shallow ensemble after #308 one gets

ValueError: Provided function should be compatible with either an energy or a force. Found a function whose output has 
shape ShapeDtypeStruct(shape=(5,), dtype=float64).