Implementation of custom Optax optimizers

[MarioAuditore]

Hello! I am wondering how one can implement their own solver with Optax API. In my case, I have my own class, which I use for riemannian gradient descent. It looks like this:

class GeometricOptimiser():
    def __init__(self, lr, maxiter, manifold, lr_reducer = None):
        
        self.lr = lr
        self.maxiter = maxiter
        self.manifold = manifold
        self.lr_reducer = lr_reducer

    def decrease_lr(self):
        if self.lr_reducer:
            return self.lr * self.lr_reducer
        else:
            return self.lr

class GradientDescent(GeometricOptimiser):
    
    @partial(jit, static_argnums=(0,))
    def step(self, param, euclid_grad):
        
        # Tangent projection for Riemannian gradient
        riem_grad = self.manifold.project(param, euclid_grad)
        # Update param
        param_updated = self.manifold.retract(param, -self.lr * riem_grad)
        # Update learning rate
        self.lr = self.decrease_lr()
        # Return result
        return param_updated

I know how to perform optimization with this class in pure JAX, but I want to wrap it in something, so it will be an Optax optimizer. Any suggestions?