Merge branch 'main' of github.com:tensorlayer/TensorLayerX into main

Author: Laicheng0830

docs/modules/ops.rst

@@ -113,11 +113,14 @@ API - Operations
    set_seed
    is_tensor
    tensor_scatter_nd_update
+   scatter_update
    diag
    mask_select
    eye
    einsum
    set_device
+   get_device
+   to_device
 
 TensorLayerX Tensor Operations
 --------------------------------
@@ -550,6 +553,10 @@ tensor_scatter_nd_update
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 .. autofunction:: tensor_scatter_nd_update
 
+scatter_update
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. autofunction:: scatter_update
+
 diag
 ^^^^^^^^^^^^^^^^^^^^^^^
 .. autofunction:: diag
@@ -568,4 +575,12 @@ einsum
 
 set_device
 ^^^^^^^^^^^^^^^^^^^^^^^
-.. autofunction:: set_device
+.. autofunction:: set_device
+
+get_device
+^^^^^^^^^^^^^^^^^^^^^^^
+.. autofunction:: get_device
+
+to_device
+^^^^^^^^^^^^^^^^^^^^^^^
+.. autofunction:: to_device

tensorlayerx/backend/__init__.py

@@ -149,6 +149,9 @@
 from .ops import eye
 from .ops import einsum
 from .ops import set_device
+from .ops import get_device
+from .ops import scatter_update
+from .ops import to_device
 # dtype
 from .ops import (
     DType, float16, float32, float64, int8, int16, int32, int64, uint8, uint16, uint32, uint64, bool, complex64,

tensorlayerx/backend/ops/__init__.py

@@ -201,6 +201,9 @@
 from .load_backend import eye
 from .load_backend import einsum
 from .load_backend import set_device
+from .load_backend import get_device
+from .load_backend import scatter_update
+from .load_backend import to_device
 # 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

@@ -68,7 +68,7 @@ def set_context(**kwargs):
 
 
 def get_tensor_shape(x):
-    return list(P.Shape()(x))
+    return list(x.shape)
 
 
 # initializers
@@ -1844,4 +1844,24 @@ def set_device(device = 'GPU', id = 0):
     if device not in ['GPU', 'CPU', 'Ascend']:
         raise ValueError ("In mindspore, only support 'CPU', 'GPU' and 'Ascend'.")
     ms.context.set_context(device_target=device)
-    ms.context.set_context(device_id = id)
+    ms.context.set_context(device_id = id)
+
+def scatter_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)
+    shape = indices.shape
+    indices = ms.ops.reshape(indices, (shape[0], 1))
+    op = ms.ops.TensorScatterUpdate()
+    return op(tensor, indices, updates)
+
+def get_device():
+    device = ms.context.get_context("device_target")
+    id = ms.context.get_context("device_id")
+    device = device + ":" +str(id)
+    return device
+
+def to_device(tensor, device = 'GPU', id = 0):
+
+    return tensor

tensorlayerx/backend/ops/paddle_backend.py

@@ -83,7 +83,7 @@ def set_context(**kwargs):
 
 
 def get_tensor_shape(x):
-    return pd.shape(x)
+    return list(x.shape)
 
 
 # initializers
@@ -1211,7 +1211,10 @@ def floor(x):
 
 
 def gather(params, indices, axis=None):
-
+    if axis < 0:
+        axis = len(params.shape) + axis
+    if axis is None:
+        axis = 0
     return pd.gather(params, indices, axis)
 
 
@@ -1859,6 +1862,21 @@ def __call__(self, *args):
 
 def set_device(device = 'GPU', id = 0):
     device = device.lower()
-    if device == 'GPU':
+    if device == 'gpu':
         device = device + ':' + str(id)
-    paddle.device.set_device(device)
+    paddle.device.set_device(device)
+
+def scatter_update(tensor, indices, updates):
+
+    return pd.scatter(tensor, indices, updates)
+
+def get_device():
+
+    return paddle.device.get_device()
+
+def to_device(tensor, device = 'GPU', id = 0):
+    device = device.upper()
+    if device == 'GPU':
+        return paddle.to_tensor(tensor, place=paddle.CUDAPlace(id))
+    if device == 'CPU':
+        return paddle.to_tensor(tensor, place=paddle.CPUPlace())

tensorlayerx/backend/ops/tensorflow_backend.py

@@ -3890,4 +3890,75 @@ def set_device(device = 'GPU', id = 0):
                     tf.config.experimental.set_memory_growth(gpu, True)
                 tf.config.experimental.set_visible_devices(gpus[id], 'GPU')
             except RuntimeError as e:
-                print(e)
+                print(e)
+
+def scatter_update(tensor, indices, updates):
+    """Applies sparse updates to a variable
+
+    Parameters
+    ----------
+    tensor : Tensor
+        A Tensor. The dim of tensor must be 1.
+    indices : Tensor
+        Indices into the tensor.
+    updates : Tensor
+        Updated values
+
+    Returns
+    -------
+        Tensor after updated.
+
+    Examples
+    ---------
+    >>> import tensorlayerx as tlx
+    >>> x = tlx.ops.ones((5,))
+    >>> indices = tlx.ops.convert_to_tensor([0,  4,  2])
+    >>> updates = tlx.ops.convert_to_tensor([1., 4., 7.])
+    >>> res = tlx.ops.scatter_update(x, indices, updates)
+    >>> [1. 1. 7. 1. 4.]
+    """
+    shape = indices.shape
+    indices = tf.reshape(indices, shape = (shape[0], 1))
+    return tf.tensor_scatter_nd_update(tensor, indices, updates)
+
+def get_device():
+    """This function can get the specified global device.
+
+    Returns
+    -------
+        The global device.
+
+    Examples
+    ---------
+    >>> import tensorlayerx as tlx
+    >>> x = tlx.ops.get_device()
+    >>> "CPU"
+    """
+    device = tf.config.experimental.get_visible_devices('GPU')
+    if len(device) == 0:
+        device = tf.config.experimental.get_visible_devices('CPU')
+    return device
+
+def to_device(tensor, device = 'GPU', id = 0):
+    """Returns a copy of Tensor in specified device.
+
+    Parameters
+    ----------
+    tensor : Tensor
+        A tensor.
+    device : str
+        The specified device. Support 'GPU' and 'CPU'. Default is 'GPU'.
+    id : int
+        The id of specified device. Default is 0.
+
+
+    Examples
+    ---------
+    >>> import tensorlayerx as tlx
+    >>> x = tlx.ops.ones((5,))
+    >>> x = tlx.ops.to_device(x, device="GPU", id=0)
+    """
+    if device is None:
+        return tensor
+    with tf.device("/" + device.upper()+':'+str(id)):
+        return tf.identity(tensor)

tensorlayerx/backend/ops/torch_backend.py

@@ -1639,7 +1639,7 @@ def is_tensor(x):
 
 def tensor_scatter_nd_update(tensor, indices, updates):
     tensor = torch.tensor(tensor)
-    indices = torch.tensor(indices)
+    indices = torch.tensor(indices, dtype=torch.long)
     updates = torch.tensor(updates)
     indices = torch.flatten(indices)
     tensor[indices] = updates
@@ -1666,6 +1666,8 @@ def mask_select(x, mask, axis = 0):
         return x[:,:,:, mask]
 
 def eye(n, m=None, dtype=None):
+    if m is None:
+        m = n
     return torch.eye(n = n, m = m, dtype =dtype)
 
 
@@ -1684,4 +1686,27 @@ def __call__(self, *args):
 def set_device(device = 'GPU', id = 0):
     if device == 'GPU':
         torch.set_default_tensor_type('torch.cuda.FloatTensor')
-        torch.cuda.set_device(id)
+        torch.cuda.set_device(id)
+
+def scatter_update(tensor, indices, updates):
+    tensor = torch.tensor(tensor)
+    indices = torch.tensor(indices, dtype=torch.long)
+    updates = torch.tensor(updates)
+    tensor[indices] = updates
+    return tensor
+
+def get_device():
+    try:
+        id = torch.cuda.current_device()
+        device = 'GPU:' + str(id)
+        return device
+    except:
+        device = 'CPU'
+        return device
+
+def to_device(tensor, device='GPU', id=0):
+    device = device.lower()
+    if device == 'gpu':
+        device = 'cuda' + ':' + str(id)
+    tensor = tensor.detach().to(device)
+    return tensor