update ops

Author: hanjr92

docs/modules/ops.rst

@@ -112,6 +112,8 @@ API - Operations
    unsorted_segment_max
    set_seed
    is_tensor
+   tensor_scatter_nd_update
+   diag
 
 TensorLayerX Tensor Operations
 --------------------------------
@@ -538,4 +540,12 @@ set_seed
 
 is_tensor
 ^^^^^^^^^^^^^^^^^^^^^^^
-.. autofunction:: is_tensor
+.. autofunction:: is_tensor
+
+tensor_scatter_nd_update
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. autofunction:: tensor_scatter_nd_update
+
+diag
+^^^^^^^^^^^^^^^^^^^^^^^
+.. autofunction:: diag

tensorlayerx/backend/__init__.py

@@ -143,6 +143,8 @@
 from .ops import unsorted_segment_max
 from .ops import set_seed
 from .ops import is_tensor
+from .ops import tensor_scatter_nd_update
+from .ops import diag
 # dtype
 from .ops import (
     DType, float16, float32, float64, int8, int16, int32, int64, uint8, uint16, uint32, uint64, bool, complex64,

tensorlayerx/backend/ops/__init__.py

@@ -195,6 +195,8 @@
 from .load_backend import unsorted_segment_max
 from .load_backend import set_seed
 from .load_backend import is_tensor
+from .load_backend import tensor_scatter_nd_update
+from .load_backend import diag
 # dtype
 from .load_backend import (
     DType, float16, float32, float64, int8, int16, int32, int64, uint8, uint16, uint32, uint64, bool, complex64,

tensorlayerx/backend/ops/mindspore_backend.py

@@ -1818,6 +1818,13 @@ def is_tensor(x):
     return isinstance(x, ms.Tensor)
 
 def tensor_scatter_nd_update(tensor, indices, updates):
-
+    if not isinstance(tensor,  ms.Tensor) or not isinstance(updates, ms.Tensor):
+        raise TypeError("tensor, updates should be Tensor, but got tensor type is {}, "
+                        "and updates type is {}.".format(type(tensor), type(updates)))
+    indices = ms.Tensor(indices)
     op = ms.ops.TensorScatterUpdate()
-    return op(tensor, indices, updates)
+    return op(tensor, indices, updates)
+
+def diag(input, diagonal=0):
+
+    return ms.numpy.diag(input, diagonal)

tensorlayerx/backend/ops/paddle_backend.py

@@ -1531,8 +1531,15 @@ def is_tensor(x):
 
 
 def tensor_scatter_nd_update(tensor, indices, updates):
+    tensor = paddle.to_tensor(tensor)
+    indices = paddle.to_tensor(indices)
+    updates = paddle.to_tensor(updates)
     a = pd.scatter_nd(indices, pd.ones_like(updates), tensor.shape)
     a = pd.multiply(tensor, -a)
     tensor = tensor + a
     x = pd.scatter_nd_add(tensor, indices, updates)
-    return x
+    return x
+
+def diag(input, diagonal=0):
+
+    return paddle.diag(input, diagonal)

tensorlayerx/backend/ops/tensorflow_backend.py

@@ -3605,13 +3605,52 @@ def tensor_scatter_nd_update(tensor, indices, updates):
 
     Parameters
     ----------
-    tensor
-    indices
-    updates
+    tensor : Tensor
+        tensor to update.
+    indices : list
+        indices to update.
+    updates : Tensor
+        value to apply at the indices
 
     Returns
     -------
+        updated Tensor.
 
+    Examples
+    ---------
+    >>> import tensorlayerx as tlx
+    >>> tensor = tlx.ops.ones(shape=(5,3))
+    >>> indices = [[0],[ 4],[ 2]]
+    >>> updates = tlx.ops.convert_to_tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    >>> new_tensor = tlx.ops.tensor_scatter_nd_update(tensor, indices, updates)
+    >>>  [[1. 2. 3.]
+    >>>  [1. 1. 1.]
+    >>>  [7. 8. 9.]
+    >>>  [1. 1. 1.]
+    >>>  [4. 5. 6.]]
     """
 
-    return tf.tensor_scatter_nd_update(tensor, indices, updates)
+    return tf.tensor_scatter_nd_update(tensor, indices, updates)
+
+def diag(input, diagonal=0):
+    """
+
+    Parameters
+    ----------
+    input : Tensor
+        the input tensor.
+    diagonal : int
+         the diagonal to consider. Defualt is 0.
+
+    Returns
+    -------
+        the output tensor.
+
+    Examples
+    ---------
+    >>> import tensorlayerx as tlx
+    >>> tensor = tlx.ops.convert_to_tensor([[1,2,3],[4,5,6],[7,8,9]]))
+    >>> new_tensor = tlx.ops.diag(tensor)
+    >>> [1, 5, 9]
+    """
+    return tf.experimental.numpy.diag(input, diagonal)

tensorlayerx/backend/ops/torch_backend.py

@@ -1633,9 +1633,13 @@ def is_tensor(x):
     return isinstance(x, torch.Tensor)
 
 def tensor_scatter_nd_update(tensor, indices, updates):
-
+    tensor = torch.tensor(tensor)
+    indices = torch.tensor(indices)
+    updates = torch.tensor(updates)
     indices = torch.flatten(indices)
     tensor[indices] = updates
     return tensor
 
+def diag(input, diagonal=0):
 
+    return torch.diag(input, diagonal)