Allow 1d and 2d normal fields in distributed mode#54

ASKabalan wants to merge 2 commits intomainDifferentiableUniverseInitiative/JaxPM:mainfrom

extend-normal_fieldDifferentiableUniverseInitiative/JaxPM:extend-normal_fieldCopy head branch name to clipboard