[GitHub] [incubator-tvm] siju-samuel commented on a change in pull request #6429: [Relay/TOPI] Added 'offsets' and 'alignment' attributes to MATRIX_SET_DIAG.

[GitBox <gi...@apache.org>]

siju-samuel commented on a change in pull request #6429:
URL: https://github.com/apache/incubator-tvm/pull/6429#discussion_r493740837



##########
File path: python/tvm/relay/op/transform.py
##########
@@ -1178,17 +1178,33 @@ def sparse_to_dense(sparse_indices, output_shape, sparse_values, default_value=0
     return _make.sparse_to_dense(sparse_indices, output_shape, sparse_values, default_value)
 
 
-def matrix_set_diag(data, diagonal):
+def matrix_set_diag(data, diagonal, k=0, align="RIGHT_LEFT"):
     """
-    Returns a tensor with the diagonal of input tensor replaced with the provided diagonal values.
+    Returns a tensor with the diagonals of input tensor replaced with the provided diagonal values.
 
     Parameters
     ----------
     data : relay.Expr
         Input Tensor.
+
     diagonal : relay.Expr
         Values to be filled in the diagonal.
 
+    k : int or tuple of int

Review comment:
       add `Optional` for k and align

##########
File path: tests/python/relay/test_op_level10.py
##########
@@ -520,17 +518,18 @@ def _verify(input_shape, dtype):
         func = relay.Function([input, diagonal], out)
         input_np = np.random.randint(-100, 100, size=input_shape).astype(dtype)
         diagonal_np = np.random.randint(-100, 100, size=diagonal_shape).astype(dtype)
-        out_np = tvm.topi.testing.matrix_set_diag(input_np, diagonal_np)
+        out_np = tvm.topi.testing.matrix_set_diag(input_np, diagonal_np, k, align)
 
         for target, ctx in tvm.testing.enabled_targets():
             for kind in ["graph", "debug"]:
                 intrp = relay.create_executor(kind, ctx=ctx, target=target)
                 out_relay = intrp.evaluate(func)(input_np, diagonal_np)
                 tvm.testing.assert_allclose(out_relay.asnumpy(), out_np)
 
-    _verify((2, 2), "float32")
-    _verify((4, 3, 3), "int32")
-    _verify((2, 3, 4), "float32")
+    _verify((2, 2), (2,), "float32")
+    _verify((4, 3, 3), (4, 3), "int32")
+    _verify((2, 3, 4), (2, 3), "float32", 1)
+    _verify((2, 3, 4), (2, 4, 3), "int32", (-1, 2), "LEFT_RIGHT")

Review comment:
       add testcases for  "LEFT_LEFT" & "RIGHT_RIGHT" 




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org