如何定义函数以构造具有给定形状的特定类型的张量

Question

如果张量的项等于索引平方和，则称其为i型。 我想要一个函数来构造具有任意形状的i型张量。可以使用numpy完成，如下所示：

import numpy as np
def i_tensor(shape = [2,2,2]):
    t = np.zeros(shape)
    it = np.nditer(t, flags=['multi_index'])
    while not it.finished:
        ind = it.multi_index
        t[ind] = np.sum(np.array(ind))
        it.iternext()
    return t

测试1：

i_tensor([3,3])

结果1：

array([[0., 1., 2.],
       [1., 2., 3.],
       [2., 3., 4.]])

Test2：

i_tensor([2,2,2])

结果2：

array([[[0., 1.],
        [1., 2.]],

       [[1., 2.],
        [2., 3.]]])

我可以使用与火炬张量相同的功能吗？我看不到nditer的张量替换。