Structures for incorporating linear functions into a nonlinear optimization problem

I'm working on a problem which naturally involves both linear and nonlinear operations, and I'd like some help understanding the best way to combine these into a neural network framework. To be more precise: suppose I have a matrix $$A$$, and some basis of matrices which it can be decomposed into, i.e.,

$$A = \sum_{j}a_{j}A_{j}$$

for some real numbers $$a_{j}$$ and some pre-specified basis matrices $$A_{j}$$. Once I have this matrix $$A$$, I want to apply a function $$f$$ to it to obtain $$f(A) = X$$. Note that in this context I mean actually applying $$f$$ to the full matrix $$A$$: it is not elementwise applications of $$f$$, i.e., $$f(A)$$ is defined by the Taylor series of $$f$$, and $$f$$ is a nonlinear function. (In case it matters, the function I have in mind is the matrix exponential, but I don't think this should make much of a difference in the statement of the problem.)

Now, I also have some objective matrix $$Y$$. My goal is to find $$A$$ such that $$X$$ approximates $$Y$$ as closely as possible. The key here is that $$f$$ is a nonlinear function, but I want to restrict $$A$$ to be a linear combination of the $$A_{j}$$'s. I don't know how to incorporate that constraint naturally in a machine learning context, but it's obviously desirable to do so since I would like to do gradient descent over both the $$a_{j}$$'s and the parameters which enforce the function $$f$$ to find the desired matrix $$X$$. Any guidance, broadly speaking about what neural network architectures are best to tackle this kind of problem would be very helpful.