I guess it can be implemented in the 'sampler' part. When solving an actual PDE, project the output of the AI onto a space that preserves the invariants.