Gradcat implementation from nest

[abdohelmy]

I was trying to implement Gradcat from nest paper but to no avail, can't seem to make it work.
Did anyone try to implement it before and got it to work, or even have the JAX code of it from the original repo, that would be an extremely huge help.

Thanks.