mars.tensor.diag#

mars.tensor.diag(v, k=0, sparse=None, gpu=None, chunk_size=None)[来源]#

提取对角线或构造对角张量。

如果您使用这个函数提取对角线并希望写入结果张量，请参阅有关mt.diagonal的更详细文档

Parameters

  • v (array_like) – 如果 v 是一个 2-D 张量，返回其 k -th 对角线。 如果 v 是一个 1-D 张量，返回一个 2-D 张量，在 k -th 对角线上包含 v。

  • k (int, 可选) – 相关的对角线。默认值为0。使用 k>0 表示主对角线以上的对角线，使用 k<0 表示主对角线以下的对角线。

  • 稀疏 (布尔值, 可选) – 如果为真则创建稀疏张量，默认为假

  • gpu (bool, 可选) – 如果为True，则在GPU上分配张量，默认为False

  • chunk_size (int 或 tuple 的 int 或 tuple 的 ints, 可选) – 每个维度上所需的块大小

Returns

out – 提取的对角线或构造的对角线张量。

Return type

张量

另请参阅

diagonal

返回指定的对角线。

diagflat

创建一个以扁平化输入为对角线的二维数组。

trace

沿对角线求和。

triu

张量的上三角部分。

tril

张量的下三角。

示例

>>> import mars.tensor as mt

>>> x = mt.arange(9).reshape((3,3))
>>> x.execute()
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

>>> mt.diag(x).execute()
array([0, 4, 8])
>>> mt.diag(x, k=1).execute()
array([1, 5])
>>> mt.diag(x, k=-1).execute()
array([3, 7])

>>> mt.diag(mt.diag(x)).execute()
array([[0, 0, 0],
       [0, 4, 0],
       [0, 0, 8]])