If you define a1 and a2 as tensors with requires_grad=True or as nn.Parameters, Autograd will compute their gradients in the backward pass.
I’m not sure how the "gradient of M wrt aX" would be calculated in your example, but if you compute a loss using M, you would be able to compute the gradient of this loss function with respect to the parameters.
I’m not familiar with your use case, but note that optimizing these parameters (aX) could just push them to negative values, which would then “minimize” the loss.