It depends what you do after, maybe if you can do it in a smooth way, but that doesn’t seem to be differentiable.
For example, if the output of your neural net is (0.5, 1.2, 1.6) and you convert this to (0, 1, 2), then this function is clearly not differentiable.
Everything depends on what you do with the matrix object after.
If you want to check that M(i, j, k) is of the correct value, for example, and you know that your matrix is smooth / almost continuous (your values are very similar in a local area) you can approximate a gradient, by writing a new autograd function that will take your three neural net outputs (x, y, z) and return the matrix output (forward).
You then have to write the backward function as well, which would basically return the local change of value of M for each direction (x, y, z).
Here’s the doc to write an autograd function:
http://pytorch.org/docs/master/notes/extending.html#extending-torch-autograd 51