[nnx] Fix State.__sub__

Commit

1 year ago

[nnx] Fix State.__sub__ Calculates the diff in a flattened representation according to the original behavior.