Does anybody know if something like this could be implemented using a library like Jax?

It's a bit outside my area, but I've seen a few encryption, or maybe signature breaking tools that somehow managed to apply some very cool linear algebra to guess at the bits directly using the gradient of the encryption function (can't find a nice link now)