code: def Kdiag(self, X):
          
          Kdiag = np.zeros(X.shape[0])
          ly=1/self.lengthscale_Y
          lu=np.sqrt(3)/self.lengthscale_U
  
          Vu = self.variance_U
          Vy=self.variance_Y
  
          k1 = (2*lu+ly)/(lu+ly)**2
          k2 = (ly-2*lu + 2*lu-ly ) / (ly-lu)**2
          k3 = 1/(lu+ly) + (lu)/(lu+ly)**2
  
          slices = index_to_slices(X[:,-1])
  
          for i, ss1 in enumerate(slices):
              for s1 in ss1:
                  if i==0:
                      Kdiag[s1]+= self.variance_U + self.ubias
                  elif i==1:
                      Kdiag[s1]+= Vu*Vy*(k1+k2+k3)
                  else:
                      raise ValueError("invalid input/output index")
          #Kdiag[slices[0][0]]+= self.variance_U   #matern32 diag
          #Kdiag[slices[1][0]]+= self.variance_U*self.variance_Y*(k1+k2+k3)  #  diag
          return Kdiag
docs: Compute the diagonal of the covariance matrix associated to X.