Optimizing a Single Channel of UV Texture in Mitsuba 3

[prk9009]

I want to: Create a scene with base_color as a differentiable UV texture, only optimize the red channel of this texture, while maintaining the computation graph for proper gradient flow

What I have tried: create a 2d array as a latent variable(representing red channel in UV space), apply a transformation where i create a new texture with G, B values from the orginal base_color texture map, and copy the red channel values to the new texture, but this is breaking the computation graph.

Question: Is there a proper way in Mitsuba 3 to optimize just one channel of a texture while maintaining the gradient flow through the computation graph? Perhaps using masks or some other technique that preserves differentiability?
Any examples or guidance would be greatly appreciated!