Fixing gradient regression in some case introduced by strictness annotation (and unsafe array indexing)