code: def trilu(data, k, upper):
      """
      Given a 2-D matrix or batches of 2-D matrices, returns the
      upper or lower triangular part of the tensor.
  
      Parameters
      ----------
      data: tvm.te.Tensor
          The tensor that trilu will be applied to. Must be either
          a 2D matrix or a tensor of batches of 2D matrices.
  
      k: tvm.te.Tensor
          The number of diagonals above or below the main diagonal
          to exclude or include.
  
      upper: bool
          If True, only upper triangular values of input are kept,
          if False, the lower triangular values are kept.
  
  
      Returns
      -------
      ret : relay.Expr
          The new tensor with appropriate diagonals set to zero.
  
      Examples
      --------
      .. code-block:: python
  
          x = [[0, 1, 2],
               [3, 4, 5],
               [6, 7, 8]]
  
          relay.trilu(x, True, 0) =
              [[0, 1, 2],
               [0, 4, 5],
               [0, 0, 8]]
      """
      # Make sure datatype is consistent.
      if k.dtype != "int32":
          k = tvm.tir.Cast("int32", k)
  
      # Check either above or below diagonal depending on upper.
      check_op = tvm.tir.GE
      if upper:
          check_op = tvm.tir.LE
  
      def _apply_trilu(*indices):
          row_index = indices[-2]
          col_index = indices[-1]
          # promote row & col indices
          if row_index.dtype != col_index.dtype:
              target_type = (col_index + row_index).dtype
              if row_index.dtype != target_type:
                  row_index = tvm.tir.Cast(target_type, row_index)
              else:
                  col_index = tvm.tir.Cast(target_type, col_index)
          other_indices = indices[:-2]
          check_position = check_op(row_index, col_index - k)
          value = data(*other_indices, row_index, col_index)
          return tvm.tir.Select(
              check_position, value, tvm.tir.const(0, data.dtype)
          )
  
      return te.compute(
          data.shape, _apply_trilu, name="trilu", tag=topi.tag.ELEMWISE
      )
  
imports: from __future__ import absolute_import
  import tvm
  from tvm import te, topi
  from tvm.te import hybrid
  from .None import cpp, tag
  from .utils import const_vector, make_idx, within_index
label_desc: Use image node_id 35 for calling a global function with example usage: trilu(data, k, upper) and returns: te, tvm