Hello there! Thanks for the reply!
I spoke with a colleague who suggested a much better solution for the point-to-image mapping. Firstly, I should say, I don’t know the correct terminology. I believe a 3D point is in the form [x, y, z, w] and a bitmap is of the form W x H x C where C is usually the colour depth / channels in this case, one channel. I thought I’d repeat it here because I’ve seen this question pop up a few times before (whilst I was looking for it) and I don’t think it’s been answered.
If one has a point and you want to see how it appears on the screen AND use this screen image as part of a loss function if needs to be differentiatable all the way through. To do this I’ve found the following:

Take your 3D point and pass it through your model -> view -> perspective -> W/divide -> ndc_to_screen matrices. This gives (Xs, Ys)
call torch.narrow() to remove z and w coordinates. This leaves you with a 2D screen coordinate
Create two matrices (2D tensors here) called XS and YS. These will be, I believe they are called, index matrices. They have the form (W,H). XS looks like:
[0,1,2,3… w]
[0,1,2,3…w]
.
.
[h]

You can call xs.permutate([1,0]) to get the ys matrix.

The clever bit. Expand your point (Xs, Ys) using expand_as(xs) and expand_as(ys) to give xe and ye.
Subtract xs from xe and the same for ys and ye. You now have two matrices that correspond to the offsets from the point you passed in. The smallest value will be 0 and that’s where your point lives in the final image.
use any function you like to combine these two matrices into a final image. One could invert, multiply and normalise for points, or use a gaussian function.

All of that is differentiable and can be used as as cost function against another image (least I hope so!)
I’ve improved the code which makes asking questions a bit easier.
The main problem is the chain of gradients and backward step. Ideally, I want to have just two variables - Xrot and Zrot. For some reason, I can only get the rot_mat to update and it’s very wrong. Not only that I have to retain the graph and retain the gradient on the rot_mat. None of this is ideal.
def new_gauss(ex, ey, xs, ys, sigma):
  return torch.clamp( torch.sum(\
      torch.exp(-((ex - xs)**2 + (ey-ys)**2)/(2*sigma**2)), dim=0),\
      0.0,1.0)
  
def test(points, rot_mat, mask) :  
  size = (128, 128) # N x C x H x W apparently
  base = test_compare()
  learning_rate = 0.1
  sigma = 2.5

  x_rot = torch.tensor([math.radians(0)], dtype=torch.float32, requires_grad=True)
  y_rot = torch.tensor([0.0], dtype=torch.float32, requires_grad=True)
  z_rot = torch.tensor([math.radians(0)], dtype=torch.float32, requires_grad=True)
  
  # Create a square array of X positions 
  numbers = list(range(0,size[0]))
  square = [ numbers for x in numbers]
  cube = []
  for i in range(0,points.shape[0]):
    cube.append(square)
  xs = torch.tensor(cube, dtype = torch.float32)
  # Now move it around to get the equivalent Y positions
  # These two matrices will be used with our Point coord
  ys = xs.permute([0, 2, 1])
  ndc_mat = gen_ndc(size)
  rot_mat = gen_rot(x_rot, y_rot, z_rot)
  trans_mat = gen_trans(0.0, 0.0, 2.0)
  proj_mat = gen_perspective(math.radians(90), 1.0, 1.0, 10.0)

  for i in range(100): 
    rot_mat.retain_grad() # Need this because it's not a leaf node?!?!
    model_mat = torch.matmul(trans_mat, rot_mat)
    o = torch.matmul(model_mat, points)
    q = torch.matmul(proj_mat, o)
    # Divide through by W seems to work with a handy mask and sum
    w = q * mask
    w = torch.sum(w, 1, keepdim=True)
    r = q / w
    s = torch.matmul(ndc_mat, r)
    t = s.narrow(1, 0, 2)

    px = t.narrow(1, 0, 1)
    py = t.narrow(1, 1, 1)

    ex = px.expand(points.shape[0], xs.shape[1], xs.shape[2])
    ey = py.expand(points.shape[0], ys.shape[1], ys.shape[2])

    model = new_gauss(ex, ey, xs, ys, sigma)
    save_image(model)
    loss = ((base-model)**2).sum()
    loss.backward(retain_graph=True)
    
    with torch.no_grad():
      #x_rot -= learning_rate * x_rot.grad
      #y_rot -= learning_rate * y_rot.grad
      #z_rot -= learning_rate * z_rot.grad
      rot_mat -= learning_rate * rot_mat.grad
      rot_mat.grad.zero_()
      splatted = splat(points, model_mat, proj_mat, ndc_mat, size)
      img = Image.fromarray(np.uint8(splatted.detach().numpy() * 255))
      img.save("torch" + str(i).zfill(3) + ".jpg", "JPEG")

  return rot_mat


The result of the intitial step is pretty good, which shows the conversion into a set of pytorch matrix functions is working. Not much to look at but it is indeed, a torus edge on.

So the main question I have - if I’m getting the gradients to flow back, why do I need retain graph and is there a reason why parts of the rotation matrix that shouldn’t be messed with are being altered? Is there a way to only be concerned with X and Z rotation angles here?
Cheers
Ben