public INDArray setDiag(INDArray in, INDArray diag) {
    NDValidation.validateNumerical("setDiag", "in", in);
    NDValidation.validateNumerical("setDiag", "diag", diag);
    return Nd4j.exec(new org.nd4j.linalg.api.ops.impl.transforms.custom.MatrixSetDiag(in, diag))[0];
  }