Constraints should take dim argument

[GeoffNN]

... to specify which dimensions are batch dims, and which are data dims.