Any fixed point solvers for convolution operation JAx.

[SubhositRay]

I tried to use the flatten after the convolution operation so I can use Andersson accelaration or other fixed point layer. But Jax gives jit and gradient the way it works gives some error. Flattening and reshaping fails during gradient descent, even though the function runs pretty well when jit and grad is not called on it. Any ideas?!

[jakevdp]

Could you edit your question to add a minimal reproducible example of the issue you're seeing? I'm having trouble understanding from your description what code you might be running.

[SubhositRay]

I am passing this function to my andersson solver:
def cell_module(params, img):

img = jnp.array(img)
img = img.reshape((1,16,40,40))
img = 1 *(lax.conv(lax.conv(img, params[0], (1,1), 'SAME') + params[1], params[2], (1,1), 'SAME')+img)
return img.flatten()

The andersson solver is from the implicit layer tutorial : http://implicit-layers-tutorial.org/implicit_functions/

Then I am running the anderson solver on the convolution of image given as:
def ca_loss(params, img, target):
img = fixed_point_layer(anderson_solver, cell_module, params, img)
img = img.reshape((1,16,40,40))
loss = jnp.mean((jnp.transpose(target, axes = (2,0,1))-img[0,:4,...])**2, axis=(1,2))
return loss.mean()

@jit
def update_params(img, params, target):
grads = grad(ca_loss)(params, img, target)
return params-step_size*grads

Now when I am trying update params in a loop I get this error

Let me know what you think is going on~

[shoyer]

JAXopt has support for fixed-point solvers, I recommend taking a look at their documentation.

[SubhositRay]

The implementation doesn't help me. 1st I am looking to implement things from scratch. Also, I tried looking into the documentation. I am trying to compute loss between target and the fixedpointlayer output (which is my implicit layer). So, taking a gradient directly on the fixedpointlayer output.params does not help me as shown in the code implementation.