I used numpy’s triu_indices to set things to 0 above the diagonal (see below). An alternative can be you could also use a mask.

  
      github.com
  
  
    t-vi/candlegp/blob/59e03ce263a1ccb5972c4eca80617e3698db049a/candlegp/conditionals.py#L21 4

# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.




import torch
from torch.autograd import Variable
import numpy


def batch_tril(A):
B = A.clone()
ii,jj = numpy.triu_indices(B.size(-2), k=1, m=B.size(-1))
B[...,ii,jj] = 0
return B


def batch_diag(A):
ii,jj = numpy.diag_indices(min(A.size(-2),A.size(-1)))
return A[...,ii,jj]







  
  
    
    
  
  


Best regards
Thomas