New set mode

[chapman20j]

What does this PR do?

Add nnx.set_mode method. Add tests for nnx.set_mode. Updates failing nnx.Dropout tests.

Fixes # (issue)

Checklist

• [ x ] This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other
  checks if that's the case).
• This change is discussed in a Github issue/
  discussion (please add a
  link).
• [ x ] The documentation and docstrings adhere to the
  documentation guidelines.
• [ x ] This change includes necessary high-coverage tests.
  (No quality testing = no merge!)