Obtaining the diagonal of Jacobians inside a piecewise function

I have a piecewise function that has in each case an N number of functions that take N arguments. I want to retrieve the diagonals of the Jacobian, or simply the element-wise first derivative of the N functions and N arguments.

Example code:

i[a_,b_]:=Piecewise[
{{{f[a,b],g[a,b]},a>b},
{{j[a,b],k[a,b]},2a<b}},
{l[a,b],m[a,b]}
]
D[i[a,b],{{a,b}}]


D can handle piecewise functions and provides the full Jacobians inside of the piecewise function.

How do I obtain a piecewise function that contains the diagonals of the Jacobian instead of the full Jacobian?

Can I rewrite this so Mathematica does not calculate the full Jacobian, which when N is large results in wasteful unnecessary computation?

EDIT In looking at MichealE2's helpful answer to this question, I was trying to think of a more general way to apply functions to expressions inside piecewise functions and found this kglr piecewise function mapping solution. Reproduced and and applied with MichaelE2's MapThread approach here:

pwMap[f_] := MapAt[f,#,{{1, All, 1},{2}}]&;


This is nice because it should be a completely general way to apply functions to expressions inside piecewise functions.

Instead of creating the whole Jacobian and then taking just the diagonal, we can just create the elements on their own with Inner. And we pick out the functions in Piecewise with the pattern funcs_List /; Depth[funcs] == 3 and replace them with the diagonal of their associated Jacobian.

i[x, y] /. funcs_List /; Depth[funcs] == 3 :> Inner[D, funcs, {x, y}, List]


$$\begin{cases} \left\{f^{(1,0)}(x,y),g^{(0,1)}(x,y)\right\} & x>y \\ \left\{j^{(1,0)}(x,y),k^{(0,1)}(x,y)\right\} & 2 x

First Q:

PiecewiseExpand@ Quiet@ Diagonal[D[i[a, b], {{a, b}}]]


Second Q:

diagJac[f_, vars_?VectorQ] := Module[{dmap},

• Thanks for this! Leveraging PiecewiseExpand to expand a function that is not evaluating(?) is clever. Though in your first solution I am not sure if suppressing the warnings would have unintended consequences? I was thinking a better way would be to have a undefined function applied and then pushed in with PiecewiseExpand, but actually found another solution in creating it (see question edit). Jan 4, 2020 at 17:43
• @JohnSmith I usually don't suppress the warnings for that very reason, but sometimes they irritate other users. Probably because the message makes them worry if there is an actual error. In this case, I don't understand why the message is issued. For instance PiecewiseExpand@Eigenvalues[D[i[a, b], {{a, b}}]] gives no message because Eigenvalues[m] will patiently wait until m is a matrix. (Eigenvalues[RandomReal[1, {3, 2}]] will give a message, because a matrix that is not square seems like an error to me.) Jan 4, 2020 at 18:25