optim: MARS

[stockeh]

Added MARS and common file for repeated ops (e.g., newton_schulz)

There is a function to set the previous gradient:

def set_last_grad(self, gradients: dict):
    """Set the last gradient for each parameter"""
    assert self._initialized, "Must be initialized before setting"
    if not self.is_approx:

        def update_last(gradient, state):
            state["last_grad"] = gradient

        tree_map(update_last, gradients, self.state)

however, this was tested extensively. The authors state in their paper p8 that using MARS-approx with is_approx = True is recommended.

Before submitting

• I've read and followed all steps in the Making a pull request
  section of the CONTRIBUTING docs.
• I've updated or added any relevant docstrings following the syntax described in the
  Writing docstrings section of the CONTRIBUTING docs.
• If this PR fixes a bug, I've added a test that will fail without my fix.
• If this PR adds a new feature, I've added tests that sufficiently cover my new functionality.