Given a convex function $g(x):\mathbb{R}^n\rightarrow \mathbb{R}$, the proximal operator of $g$ is defined as
$P_g(x)=\underset{u}{\arg\min}\quad \frac{1}{2}||x-u||_2^2+g(u)$.
Since $g(x)$ is convex, the proximal is a singleton, i.e., there is a unique minimizer $u$. Thus, the proximal operator is a vector function $P_g:\mathbb{R}^n\rightarrow\mathbb{R}^n$.
I'm trying to find an expression for the Jacobian (or weak Jacobian) of $P_g$. In the case where $g$ is differentiable, I believe I can find the Jacobian using the implicit function theorem. However, I'm interested in the general case where $g$ is convex but not differentiable.
As first try, I tried to look on the Moreau Envelope of $g$: $M_g(x)=\underset{u}{\min}\quad \frac{1}{2}||x-u||_2^2+g(u)$.
This function is differentiable and its gradient is given by
$\nabla M_g(x) = x-P_g(x)$ .
Therefore, if I could compute the Hessian $H_g$ of $M_g$, I'll get
$H_g=I-J_g$
where $J_g$ is the desired Jacobian.
However, I'm not sure under which conditions the Hessian exists (even in the weak sense) and what the expression for it? Moreover, maybe there is another way to approach this.
A good reference to start on this is probably - Generalized Hessian Properties of Regularized Non Smooth Functions.
If you look at Thm. 3.8, it gives the same relationship as you derived between Hessian of the Moreau envelope and Jacobian of prox. Furthermore, it gives a necessary and sufficient condition for the Hessian of the Moreau envelope to exist, that is the second-order epi-derivative of your $g$ being generalized quadratic.
And this is probably where it is better to start digging into the yellow book: it looks like some version of it is available - R. Tyrrell Rockafellar, Roger J-B Wets - Variational Analysis.
Hope this helps!