Question - custom losses.

[frank-ccc]

Dear All,
Thank you so much for writing such a great code.
I am sorry to ask this question here; I do hope this is the right place to post questions
I was wondering if there are any examples of custom losses for inverse problems. In particular, I would like to

• ask how custom loss could be implemented (I am interested in a loss like np.sum((A*x-b)**2))
• suggest providing an example of custom loss in the colab notebooks
  Thank you so much

[romanngg]

Thanks for the question! Could you please provide more details, in particular clarify the dimensionality of A, namely is this a constant matrix of shape output_size x output_size that you apply to each training sample, or is it an (output_size x trainig_set_size) x (output_size x trainig_set_size) matrix, or something else?

If it's the former, then I think you can model it with something like stax.DotGeneral(rhs=A) layer at the top (and b being your train targets y_train), please see examples in https://neural-tangents.readthedocs.io/en/latest/_autosummary/neural_tangents.stax.DotGeneral.html. If not, please describe the problem in more detail/specific dimensions, so that I understand it better.

Also related: #55

[frank-ccc]

Dear Roman,
Thank you for the reply. A 2N\times N is a constant matrix, x is an N vector, the value of b is the target a 2N constant vector. The setting I am interested in is that of a one-shot inverse problem (in some sense it is similar to the DIP, see https://arxiv.org/abs/1711.10925). I looked at the documentation and I think Dot_general may work, I will try it.
Thank you again