Skip to content

D
DI-treetensor

OpenDILab开源决策智能平台 / DI-treetensor
11 个月 前同步成功

通知 42
Star 172
Fork 11
D
DI-treetensor

体验新版 GitCode，发现更多精彩内容 >>

提交 b9f0c0bc 编写于 作者: HansBug's avatar HansBug 😆
浏览文件
操作

dev, test(hansbug): complete add, sub, mul, div, pow, neg

上级 7c15d9f7
隐藏空白更改
内联 并排
Showing with 1017 addition and 0 deletion
test/torch/test_funcs.py
浏览文件 @ b9f0c0bc
...@@ -964,3 +964,233 @@ class TestTorchFuncs: ...@@ -964,3 +964,233 @@ class TestTorchFuncs:
'b': {'x': [[0.6225, 0.7685], 'b': {'x': [[0.6225, 0.7685],
[0.0759, 0.5622]]}, [0.0759, 0.5622]]},
})) < 1e-4).all() })) < 1e-4).all()
@choose_mark()
def test_add(self):
t1 = ttorch.add(
ttorch.tensor([1, 2, 3]),
ttorch.tensor([3, 5, 11]),
)
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([4, 7, 14])).all()
t2 = ttorch.add(
ttorch.tensor({
'a': [1, 2, 3],
'b': {'x': [[3, 5], [9, 12]]},
}),
ttorch.tensor({
'a': [3, 5, 11],
'b': {'x': [[31, -15], [13, 23]]},
})
)
assert (t2 == ttorch.tensor({
'a': [4, 7, 14],
'b': {'x': [[34, -10],
[22, 35]]},
})).all()
@choose_mark()
def test_sub(self):
t1 = ttorch.sub(
ttorch.tensor([1, 2, 3]),
ttorch.tensor([3, 5, 11]),
)
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([-2, -3, -8])).all()
t2 = ttorch.sub(
ttorch.tensor({
'a': [1, 2, 3],
'b': {'x': [[3, 5], [9, 12]]},
}),
ttorch.tensor({
'a': [3, 5, 11],
'b': {'x': [[31, -15], [13, 23]]},
})
)
assert (t2 == ttorch.tensor({
'a': [-2, -3, -8],
'b': {'x': [[-28, 20],
[-4, -11]]},
})).all()
@choose_mark()
def test_mul(self):
t1 = ttorch.mul(
ttorch.tensor([1, 2, 3]),
ttorch.tensor([3, 5, 11]),
)
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([3, 10, 33])).all()
t2 = ttorch.mul(
ttorch.tensor({
'a': [1, 2, 3],
'b': {'x': [[3, 5], [9, 12]]},
}),
ttorch.tensor({
'a': [3, 5, 11],
'b': {'x': [[31, -15], [13, 23]]},
})
)
assert (t2 == ttorch.tensor({
'a': [3, 10, 33],
'b': {'x': [[93, -75],
[117, 276]]},
})).all()
@choose_mark()
def test_div(self):
t1 = ttorch.div(ttorch.tensor([0.3810, 1.2774, -0.2972, -0.3719, 0.4637]), 0.5)
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([0.7620, 2.5548, -0.5944, -0.7438, 0.9274])).all()
t2 = ttorch.div(
ttorch.tensor([1.3119, 0.0928, 0.4158, 0.7494, 0.3870]),
ttorch.tensor([-1.7501, -1.4652, 0.1379, -1.1252, 0.0380]),
)
assert isinstance(t2, torch.Tensor)
assert (ttorch.abs(t2 - ttorch.tensor([-0.7496, -0.0633, 3.0152, -0.6660, 10.1842])) < 1e-4).all()
t3 = ttorch.div(
ttorch.tensor({
'a': [0.3810, 1.2774, -0.2972, -0.3719, 0.4637],
'b': {
'x': [1.3119, 0.0928, 0.4158, 0.7494, 0.3870],
'y': [[[1.9579, -0.0335, 0.1178],
[0.8287, 1.4520, -0.4696]],
[[-2.1659, -0.5831, 0.4080],
[0.1400, 0.8122, 0.5380]]],
},
}),
ttorch.tensor({
'a': 0.5,
'b': {
'x': [-1.7501, -1.4652, 0.1379, -1.1252, 0.0380],
'y': [[[-1.3136, 0.7785, -0.7290],
[0.6025, 0.4635, -1.1882]],
[[0.2756, -0.4483, -0.2005],
[0.9587, 1.4623, -2.8323]]],
},
}),
)
assert (ttorch.abs(t3 - ttorch.tensor({
'a': [0.7620, 2.5548, -0.5944, -0.7438, 0.9274],
'b': {
'x': [-0.7496, -0.0633, 3.0152, -0.6660, 10.1842],
'y': [[[-1.4905, -0.0430, -0.1616],
[1.3754, 3.1327, 0.3952]],
[[-7.8589, 1.3007, -2.0349],
[0.1460, 0.5554, -0.1900]]],
}
})) < 1e-4).all()
@choose_mark()
def test_pow(self):
t1 = ttorch.pow(
ttorch.tensor([4, 3, 2, 6, 2]),
ttorch.tensor([4, 2, 6, 4, 3]),
)
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([256, 9, 64, 1296, 8])).all()
t2 = ttorch.pow(
ttorch.tensor({
'a': [4, 3, 2, 6, 2],
'b': {
'x': [[3, 4, 6],
[6, 3, 5]],
'y': [[[3, 5, 5],
[5, 7, 6]],
[[4, 6, 5],
[7, 2, 7]]],
},
}),
ttorch.tensor({
'a': [4, 2, 6, 4, 3],
'b': {
'x': [[7, 4, 6],
[5, 2, 6]],
'y': [[[7, 2, 2],
[2, 3, 2]],
[[5, 2, 6],
[7, 3, 4]]],
},
}),
)
assert (t2 == ttorch.tensor({
'a': [256, 9, 64, 1296, 8],
'b': {
'x': [[2187, 256, 46656],
[7776, 9, 15625]],
'y': [[[2187, 25, 25],
[25, 343, 36]],
[[1024, 36, 15625],
[823543, 8, 2401]]],
}
})).all()
@choose_mark()
def test_neg(self):
t1 = ttorch.neg(ttorch.tensor([4, 3, 2, 6, 2]))
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([-4, -3, -2, -6, -2])).all()
t2 = ttorch.neg(ttorch.tensor({
'a': [4, 3, 2, 6, 2],
'b': {
'x': [[3, 4, 6],
[6, 3, 5]],
'y': [[[3, 5, 5],
[5, 7, 6]],
[[4, 6, 5],
[7, 2, 7]]],
},
}))
assert (t2 == ttorch.tensor({
'a': [-4, -3, -2, -6, -2],
'b': {
'x': [[-3, -4, -6],
[-6, -3, -5]],
'y': [[[-3, -5, -5],
[-5, -7, -6]],
[[-4, -6, -5],
[-7, -2, -7]]],
}
}))
@choose_mark()
def test_neg_(self):
t1 = ttorch.tensor([4, 3, 2, 6, 2])
t1r = ttorch.neg_(t1)
assert t1r is t1
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([-4, -3, -2, -6, -2])).all()
t2 = ttorch.tensor({
'a': [4, 3, 2, 6, 2],
'b': {
'x': [[3, 4, 6],
[6, 3, 5]],
'y': [[[3, 5, 5],
[5, 7, 6]],
[[4, 6, 5],
[7, 2, 7]]],
},
})
t2r = ttorch.neg_(t2)
assert t2r is t2
assert (t2 == ttorch.tensor({
'a': [-4, -3, -2, -6, -2],
'b': {
'x': [[-3, -4, -6],
[-6, -3, -5]],
'y': [[[-3, -5, -5],
[-5, -7, -6]],
[[-4, -6, -5],
[-7, -2, -7]]],
}
}))
test/torch/test_tensor.py
浏览文件 @ b9f0c0bc
...@@ -590,3 +590,383 @@ class TestTorchTensor: ...@@ -590,3 +590,383 @@ class TestTorchTensor:
'b': {'x': [[0.6225, 0.7685], 'b': {'x': [[0.6225, 0.7685],
[0.0759, 0.5622]]}, [0.0759, 0.5622]]},
})) < 1e-4).all() })) < 1e-4).all()
@choose_mark()
def test_add(self):
t1 = ttorch.tensor([1, 2, 3]).add(
ttorch.tensor([3, 5, 11]),
)
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([4, 7, 14])).all()
t2 = ttorch.tensor({
'a': [1, 2, 3],
'b': {'x': [[3, 5], [9, 12]]},
}).add(
ttorch.tensor({
'a': [3, 5, 11],
'b': {'x': [[31, -15], [13, 23]]},
})
)
assert (t2 == ttorch.tensor({
'a': [4, 7, 14],
'b': {'x': [[34, -10],
[22, 35]]},
})).all()
@choose_mark()
def test_add_(self):
t1 = ttorch.tensor([1, 2, 3])
t1r = t1.add_(ttorch.tensor([3, 5, 11]))
assert t1r is t1
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([4, 7, 14])).all()
t2 = ttorch.tensor({
'a': [1, 2, 3],
'b': {'x': [[3, 5], [9, 12]]},
})
t2r = t2.add_(ttorch.tensor({
'a': [3, 5, 11],
'b': {'x': [[31, -15], [13, 23]]},
}))
assert t2r is t2
assert (t2 == ttorch.tensor({
'a': [4, 7, 14],
'b': {'x': [[34, -10],
[22, 35]]},
})).all()
@choose_mark()
def test_sub(self):
t1 = ttorch.tensor([1, 2, 3]).sub(
ttorch.tensor([3, 5, 11]),
)
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([-2, -3, -8])).all()
t2 = ttorch.tensor({
'a': [1, 2, 3],
'b': {'x': [[3, 5], [9, 12]]},
}).sub(
ttorch.tensor({
'a': [3, 5, 11],
'b': {'x': [[31, -15], [13, 23]]},
})
)
assert (t2 == ttorch.tensor({
'a': [-2, -3, -8],
'b': {'x': [[-28, 20],
[-4, -11]]},
})).all()
@choose_mark()
def test_sub_(self):
t1 = ttorch.tensor([1, 2, 3])
t1r = t1.sub_(ttorch.tensor([3, 5, 11]))
assert t1r is t1
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([-2, -3, -8])).all()
t2 = ttorch.tensor({
'a': [1, 2, 3],
'b': {'x': [[3, 5], [9, 12]]},
})
t2r = t2.sub_(ttorch.tensor({
'a': [3, 5, 11],
'b': {'x': [[31, -15], [13, 23]]},
}))
assert t2r is t2
assert (t2 == ttorch.tensor({
'a': [-2, -3, -8],
'b': {'x': [[-28, 20],
[-4, -11]]},
})).all()
@choose_mark()
def test_mul(self):
t1 = ttorch.tensor([1, 2, 3]).mul(
ttorch.tensor([3, 5, 11]),
)
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([3, 10, 33])).all()
t2 = ttorch.tensor({
'a': [1, 2, 3],
'b': {'x': [[3, 5], [9, 12]]},
}).mul(
ttorch.tensor({
'a': [3, 5, 11],
'b': {'x': [[31, -15], [13, 23]]},
})
)
assert (t2 == ttorch.tensor({
'a': [3, 10, 33],
'b': {'x': [[93, -75],
[117, 276]]},
})).all()
@choose_mark()
def test_mul_(self):
t1 = ttorch.tensor([1, 2, 3])
t1r = t1.mul_(ttorch.tensor([3, 5, 11]))
assert t1r is t1
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([3, 10, 33])).all()
t2 = ttorch.tensor({
'a': [1, 2, 3],
'b': {'x': [[3, 5], [9, 12]]},
})
t2r = t2.mul_(ttorch.tensor({
'a': [3, 5, 11],
'b': {'x': [[31, -15], [13, 23]]},
}))
assert t2r is t2
assert (t2 == ttorch.tensor({
'a': [3, 10, 33],
'b': {'x': [[93, -75],
[117, 276]]},
})).all()
@choose_mark()
def test_div(self):
t1 = ttorch.tensor([0.3810, 1.2774, -0.2972, -0.3719, 0.4637]).div(0.5)
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([0.7620, 2.5548, -0.5944, -0.7438, 0.9274])).all()
t2 = ttorch.tensor([1.3119, 0.0928, 0.4158, 0.7494, 0.3870]).div(
ttorch.tensor([-1.7501, -1.4652, 0.1379, -1.1252, 0.0380]),
)
assert isinstance(t2, torch.Tensor)
assert (ttorch.abs(t2 - ttorch.tensor([-0.7496, -0.0633, 3.0152, -0.6660, 10.1842])) < 1e-4).all()
t3 = ttorch.tensor({
'a': [0.3810, 1.2774, -0.2972, -0.3719, 0.4637],
'b': {
'x': [1.3119, 0.0928, 0.4158, 0.7494, 0.3870],
'y': [[[1.9579, -0.0335, 0.1178],
[0.8287, 1.4520, -0.4696]],
[[-2.1659, -0.5831, 0.4080],
[0.1400, 0.8122, 0.5380]]],
},
}).div(
ttorch.tensor({
'a': 0.5,
'b': {
'x': [-1.7501, -1.4652, 0.1379, -1.1252, 0.0380],
'y': [[[-1.3136, 0.7785, -0.7290],
[0.6025, 0.4635, -1.1882]],
[[0.2756, -0.4483, -0.2005],
[0.9587, 1.4623, -2.8323]]],
},
}),
)
assert (ttorch.abs(t3 - ttorch.tensor({
'a': [0.7620, 2.5548, -0.5944, -0.7438, 0.9274],
'b': {
'x': [-0.7496, -0.0633, 3.0152, -0.6660, 10.1842],
'y': [[[-1.4905, -0.0430, -0.1616],
[1.3754, 3.1327, 0.3952]],
[[-7.8589, 1.3007, -2.0349],
[0.1460, 0.5554, -0.1900]]],
}
})) < 1e-4).all()
@choose_mark()
def test_div_(self):
t1 = ttorch.tensor([0.3810, 1.2774, -0.2972, -0.3719, 0.4637])
t1r = t1.div_(0.5)
assert t1r is t1
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([0.7620, 2.5548, -0.5944, -0.7438, 0.9274])).all()
t2 = ttorch.tensor([1.3119, 0.0928, 0.4158, 0.7494, 0.3870])
t2r = t2.div_(ttorch.tensor([-1.7501, -1.4652, 0.1379, -1.1252, 0.0380]))
assert t2r is t2
assert isinstance(t2, torch.Tensor)
assert (ttorch.abs(t2 - ttorch.tensor([-0.7496, -0.0633, 3.0152, -0.6660, 10.1842])) < 1e-4).all()
t3 = ttorch.tensor({
'a': [0.3810, 1.2774, -0.2972, -0.3719, 0.4637],
'b': {
'x': [1.3119, 0.0928, 0.4158, 0.7494, 0.3870],
'y': [[[1.9579, -0.0335, 0.1178],
[0.8287, 1.4520, -0.4696]],
[[-2.1659, -0.5831, 0.4080],
[0.1400, 0.8122, 0.5380]]],
},
})
t3r = t3.div_(ttorch.tensor({
'a': 0.5,
'b': {
'x': [-1.7501, -1.4652, 0.1379, -1.1252, 0.0380],
'y': [[[-1.3136, 0.7785, -0.7290],
[0.6025, 0.4635, -1.1882]],
[[0.2756, -0.4483, -0.2005],
[0.9587, 1.4623, -2.8323]]],
},
}))
assert t3r is t3
assert (ttorch.abs(t3 - ttorch.tensor({
'a': [0.7620, 2.5548, -0.5944, -0.7438, 0.9274],
'b': {
'x': [-0.7496, -0.0633, 3.0152, -0.6660, 10.1842],
'y': [[[-1.4905, -0.0430, -0.1616],
[1.3754, 3.1327, 0.3952]],
[[-7.8589, 1.3007, -2.0349],
[0.1460, 0.5554, -0.1900]]],
}
})) < 1e-4).all()
@choose_mark()
def test_pow(self):
t1 = ttorch.tensor([4, 3, 2, 6, 2]).pow(
ttorch.tensor([4, 2, 6, 4, 3]),
)
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([256, 9, 64, 1296, 8])).all()
t2 = ttorch.tensor({
'a': [4, 3, 2, 6, 2],
'b': {
'x': [[3, 4, 6],
[6, 3, 5]],
'y': [[[3, 5, 5],
[5, 7, 6]],
[[4, 6, 5],
[7, 2, 7]]],
},
}).pow(
ttorch.tensor({
'a': [4, 2, 6, 4, 3],
'b': {
'x': [[7, 4, 6],
[5, 2, 6]],
'y': [[[7, 2, 2],
[2, 3, 2]],
[[5, 2, 6],
[7, 3, 4]]],
},
}),
)
assert (t2 == ttorch.tensor({
'a': [256, 9, 64, 1296, 8],
'b': {
'x': [[2187, 256, 46656],
[7776, 9, 15625]],
'y': [[[2187, 25, 25],
[25, 343, 36]],
[[1024, 36, 15625],
[823543, 8, 2401]]],
}
})).all()
@choose_mark()
def test_pow_(self):
t1 = ttorch.tensor([4, 3, 2, 6, 2])
t1r = t1.pow_(ttorch.tensor([4, 2, 6, 4, 3]))
assert t1r is t1
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([256, 9, 64, 1296, 8])).all()
t2 = ttorch.tensor({
'a': [4, 3, 2, 6, 2],
'b': {
'x': [[3, 4, 6],
[6, 3, 5]],
'y': [[[3, 5, 5],
[5, 7, 6]],
[[4, 6, 5],
[7, 2, 7]]],
},
})
t2r = t2.pow_(ttorch.tensor({
'a': [4, 2, 6, 4, 3],
'b': {
'x': [[7, 4, 6],
[5, 2, 6]],
'y': [[[7, 2, 2],
[2, 3, 2]],
[[5, 2, 6],
[7, 3, 4]]],
},
}))
assert t2r is t2
assert (t2 == ttorch.tensor({
'a': [256, 9, 64, 1296, 8],
'b': {
'x': [[2187, 256, 46656],
[7776, 9, 15625]],
'y': [[[2187, 25, 25],
[25, 343, 36]],
[[1024, 36, 15625],
[823543, 8, 2401]]],
}
})).all()
@choose_mark()
def test_neg(self):
t1 = ttorch.tensor([4, 3, 2, 6, 2]).neg()
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([-4, -3, -2, -6, -2])).all()
t2 = ttorch.tensor({
'a': [4, 3, 2, 6, 2],
'b': {
'x': [[3, 4, 6],
[6, 3, 5]],
'y': [[[3, 5, 5],
[5, 7, 6]],
[[4, 6, 5],
[7, 2, 7]]],
},
}).neg()
assert (t2 == ttorch.tensor({
'a': [-4, -3, -2, -6, -2],
'b': {
'x': [[-3, -4, -6],
[-6, -3, -5]],
'y': [[[-3, -5, -5],
[-5, -7, -6]],
[[-4, -6, -5],
[-7, -2, -7]]],
}
}))
@choose_mark()
def test_neg_(self):
t1 = ttorch.tensor([4, 3, 2, 6, 2])
t1r = t1.neg_()
assert t1r is t1
assert isinstance(t1, torch.Tensor)
assert (t1 == ttorch.tensor([-4, -3, -2, -6, -2])).all()
t2 = ttorch.tensor({
'a': [4, 3, 2, 6, 2],
'b': {
'x': [[3, 4, 6],
[6, 3, 5]],
'y': [[[3, 5, 5],
[5, 7, 6]],
[[4, 6, 5],
[7, 2, 7]]],
},
})
t2r = t2.neg_()
assert t2r is t2
assert (t2 == ttorch.tensor({
'a': [-4, -3, -2, -6, -2],
'b': {
'x': [[-3, -4, -6],
[-6, -3, -5]],
'y': [[[-3, -5, -5],
[-5, -7, -6]],
[[-4, -6, -5],
[-7, -2, -7]]],
}
}))
treetensor/torch/funcs.py
浏览文件 @ b9f0c0bc
...@@ -27,6 +27,7 @@ __all__ = [ ...@@ -27,6 +27,7 @@ __all__ = [
'isfinite', 'isinf', 'isnan', 'isfinite', 'isinf', 'isnan',
'abs', 'abs_', 'clamp', 'clamp_', 'sign', 'sigmoid', 'sigmoid_', 'abs', 'abs_', 'clamp', 'clamp_', 'sign', 'sigmoid', 'sigmoid_',
'round', 'round_', 'floor', 'floor_', 'ceil', 'ceil_', 'round', 'round_', 'floor', 'floor_', 'ceil', 'ceil_',
'add', 'sub', 'mul', 'div', 'pow', 'neg', 'neg_',
] ]
func_treelize = post_process(post_process(args_mapping( func_treelize = post_process(post_process(args_mapping(
...@@ -1442,4 +1443,308 @@ def sigmoid_(input): ...@@ -1442,4 +1443,308 @@ def sigmoid_(input):
return torch.sigmoid_(input) return torch.sigmoid_(input)
# noinspection PyShadowingBuiltins
@doc_from_base()
@func_treelize()
def add(input, other, *args, **kwargs):
"""
Adds the scalar ``other`` to each element of the ``input`` input and
returns a new resulting tree tensor.
Examples::
>>> import torch
>>> import treetensor.torch as ttorch
>>> ttorch.add(
... ttorch.tensor([1, 2, 3]),
... ttorch.tensor([3, 5, 11]),
... )
tensor([ 4, 7, 14])
>>> ttorch.add(
... ttorch.tensor({
... 'a': [1, 2, 3],
... 'b': {'x': [[3, 5], [9, 12]]},
... }),
... ttorch.tensor({
... 'a': [3, 5, 11],
... 'b': {'x': [[31, -15], [13, 23]]},
... })
... )
<Tensor 0x7f11b139c710>
├── a --> tensor([ 4, 7, 14])
└── b --> <Tensor 0x7f11b139c630>
└── x --> tensor([[ 34, -10],
[ 22, 35]])
"""
return torch.add(input, other, *args, **kwargs)
# noinspection PyShadowingBuiltins
@doc_from_base()
@func_treelize()
def sub(input, other, *args, **kwargs):
"""
Subtracts ``other``, scaled by ``alpha``, from ``input``.
Examples::
>>> import torch
>>> import treetensor.torch as ttorch
>>> ttorch.sub(
... ttorch.tensor([1, 2, 3]),
... ttorch.tensor([3, 5, 11]),
... )
tensor([-2, -3, -8])
>>> ttorch.sub(
... ttorch.tensor({
... 'a': [1, 2, 3],
... 'b': {'x': [[3, 5], [9, 12]]},
... }),
... ttorch.tensor({
... 'a': [3, 5, 11],
... 'b': {'x': [[31, -15], [13, 23]]},
... })
... )
<Tensor 0x7f11b139ccc0>
├── a --> tensor([-2, -3, -8])
└── b --> <Tensor 0x7f11b139cc18>
└── x --> tensor([[-28, 20],
[ -4, -11]])
"""
return torch.sub(input, other, *args, **kwargs)
# noinspection PyShadowingBuiltins
@doc_from_base()
@func_treelize()
def mul(input, other, *args, **kwargs):
"""
Multiplies each element of the input ``input`` with the scalar ``other`` and
returns a new resulting tensor.
Examples::
>>> import torch
>>> import treetensor.torch as ttorch
>>> ttorch.mul(
... ttorch.tensor([1, 2, 3]),
... ttorch.tensor([3, 5, 11]),
... )
tensor([ 3, 10, 33])
>>> ttorch.mul(
... ttorch.tensor({
... 'a': [1, 2, 3],
... 'b': {'x': [[3, 5], [9, 12]]},
... }),
... ttorch.tensor({
... 'a': [3, 5, 11],
... 'b': {'x': [[31, -15], [13, 23]]},
... })
... )
<Tensor 0x7f11b139ca58>
├── a --> tensor([ 3, 10, 33])
└── b --> <Tensor 0x7f11b139cb00>
└── x --> tensor([[ 93, -75],
[117, 276]])
"""
return torch.mul(input, other, *args, **kwargs)
# noinspection PyShadowingBuiltins
@doc_from_base()
@func_treelize()
def div(input, other, *args, **kwargs):
"""
Divides each element of the input ``input`` by the corresponding element of ``other``.
Examples::
>>> import torch
>>> import treetensor.torch as ttorch
>>> ttorch.div(ttorch.tensor([ 0.3810, 1.2774, -0.2972, -0.3719, 0.4637]), 0.5)
tensor([ 0.7620, 2.5548, -0.5944, -0.7438, 0.9274])
>>> ttorch.div(
... ttorch.tensor([1.3119, 0.0928, 0.4158, 0.7494, 0.3870]),
... ttorch.tensor([-1.7501, -1.4652, 0.1379, -1.1252, 0.0380]),
... )
tensor([-0.7496, -0.0633, 3.0152, -0.6660, 10.1842])
>>> ttorch.div(
... ttorch.tensor({
... 'a': [ 0.3810, 1.2774, -0.2972, -0.3719, 0.4637],
... 'b': {
... 'x': [1.3119, 0.0928, 0.4158, 0.7494, 0.3870],
... 'y': [[[ 1.9579, -0.0335, 0.1178],
... [ 0.8287, 1.4520, -0.4696]],
... [[-2.1659, -0.5831, 0.4080],
... [ 0.1400, 0.8122, 0.5380]]],
... },
... }),
... ttorch.tensor({
... 'a': 0.5,
... 'b': {
... 'x': [-1.7501, -1.4652, 0.1379, -1.1252, 0.0380],
... 'y': [[[-1.3136, 0.7785, -0.7290],
... [ 0.6025, 0.4635, -1.1882]],
... [[ 0.2756, -0.4483, -0.2005],
... [ 0.9587, 1.4623, -2.8323]]],
... },
... }),
... )
<Tensor 0x7f11b139c198>
├── a --> tensor([ 0.7620, 2.5548, -0.5944, -0.7438, 0.9274])
└── b --> <Tensor 0x7f11b139c320>
├── x --> tensor([-0.7496, -0.0633, 3.0152, -0.6660, 10.1842])
└── y --> tensor([[[-1.4905, -0.0430, -0.1616],
[ 1.3754, 3.1327, 0.3952]],
[[-7.8589, 1.3007, -2.0349],
[ 0.1460, 0.5554, -0.1900]]])
"""
return torch.div(input, other, *args, **kwargs)
# noinspection PyShadowingBuiltins
@doc_from_base()
@func_treelize()
def pow(input, exponent, *args, **kwargs):
"""
Takes the power of each element in ``input`` with ``exponent`` and
returns a tensor with the result.
Examples::
>>> import torch
>>> import treetensor.torch as ttorch
>>> ttorch.pow(
... ttorch.tensor([4, 3, 2, 6, 2]),
... ttorch.tensor([4, 2, 6, 4, 3]),
... )
tensor([ 256, 9, 64, 1296, 8])
>>> ttorch.pow(
... ttorch.tensor({
... 'a': [4, 3, 2, 6, 2],
... 'b': {
... 'x': [[3, 4, 6],
... [6, 3, 5]],
... 'y': [[[3, 5, 5],
... [5, 7, 6]],
... [[4, 6, 5],
... [7, 2, 7]]],
... },
... }),
... ttorch.tensor({
... 'a': [4, 2, 6, 4, 3],
... 'b': {
... 'x': [[7, 4, 6],
... [5, 2, 6]],
... 'y': [[[7, 2, 2],
... [2, 3, 2]],
... [[5, 2, 6],
... [7, 3, 4]]],
... },
... }),
... )
<Tensor 0x7f11b13b6e48>
├── a --> tensor([ 256, 9, 64, 1296, 8])
└── b --> <Tensor 0x7f11b13b6d68>
├── x --> tensor([[ 2187, 256, 46656],
│ [ 7776, 9, 15625]])
└── y --> tensor([[[ 2187, 25, 25],
[ 25, 343, 36]],
[[ 1024, 36, 15625],
[823543, 8, 2401]]])
"""
return torch.pow(input, exponent, *args, **kwargs)
# noinspection PyShadowingBuiltins
@doc_from_base()
@func_treelize()
def neg(input, *args, **kwargs):
"""
Returns a new tensor with the negative of the elements of ``input``.
Examples::
>>> import torch
>>> import treetensor.torch as ttorch
>>> ttorch.neg(ttorch.tensor([4, 3, 2, 6, 2]))
tensor([-4, -3, -2, -6, -2])
>>> ttorch.neg(ttorch.tensor({
... 'a': [4, 3, 2, 6, 2],
... 'b': {
... 'x': [[3, 4, 6],
... [6, 3, 5]],
... 'y': [[[3, 5, 5],
... [5, 7, 6]],
... [[4, 6, 5],
... [7, 2, 7]]],
... },
... }))
<Tensor 0x7f11b13b5860>
├── a --> tensor([-4, -3, -2, -6, -2])
└── b --> <Tensor 0x7f11b13b5828>
├── x --> tensor([[-3, -4, -6],
│ [-6, -3, -5]])
└── y --> tensor([[[-3, -5, -5],
[-5, -7, -6]],
[[-4, -6, -5],
[-7, -2, -7]]])
"""
return torch.neg(input, *args, **kwargs)
# noinspection PyShadowingBuiltins
@doc_from_base()
@return_self
@func_treelize()
def neg_(input):
"""
In-place version of :func:`treetensor.torch.neg`.
Examples::
>>> import torch
>>> import treetensor.torch as ttorch
>>> t = ttorch.tensor([4, 3, 2, 6, 2])
>>> ttorch.neg_(t)
>>> t
tensor([-4, -3, -2, -6, -2])
>>> t = ttorch.tensor({
... 'a': [4, 3, 2, 6, 2],
... 'b': {
... 'x': [[3, 4, 6],
... [6, 3, 5]],
... 'y': [[[3, 5, 5],
... [5, 7, 6]],
... [[4, 6, 5],
... [7, 2, 7]]],
... },
... })
>>> ttorch.neg_(t)
>>> t
<Tensor 0x7f11b13b6fd0>
├── a --> tensor([-4, -3, -2, -6, -2])
└── b --> <Tensor 0x7f11b13b60f0>
├── x --> tensor([[-3, -4, -6],
│ [-6, -3, -5]])
└── y --> tensor([[[-3, -5, -5],
[-5, -7, -6]],
[[-4, -6, -5],
[-7, -2, -7]]])
"""
return torch.neg_(input)
sys.modules[__name__] = module_autoremove(sys.modules[__name__]) sys.modules[__name__] = module_autoremove(sys.modules[__name__])
treetensor/torch/tensor.py
浏览文件 @ b9f0c0bc
...@@ -430,3 +430,105 @@ class Tensor(Torch, metaclass=clsmeta(_to_tensor, allow_dict=True)): ...@@ -430,3 +430,105 @@ class Tensor(Torch, metaclass=clsmeta(_to_tensor, allow_dict=True)):
See :func:`treetensor.torch.round_`. See :func:`treetensor.torch.round_`.
""" """
return self.round_(*args, **kwargs) return self.round_(*args, **kwargs)
@doc_from_base()
@method_treelize()
def add(self, other, *args, **kwargs):
"""
See :func:`treetensor.torch.add`.
"""
return self.add(other, *args, **kwargs)
@doc_from_base()
@return_self
@method_treelize()
def add_(self, other, *args, **kwargs):
"""
In-place version of :meth:`Tensor.add`.
"""
return self.add_(other, *args, **kwargs)
@doc_from_base()
@method_treelize()
def sub(self, other, *args, **kwargs):
"""
See :func:`treetensor.torch.sub`.
"""
return self.sub(other, *args, **kwargs)
@doc_from_base()
@return_self
@method_treelize()
def sub_(self, other, *args, **kwargs):
"""
In-place version of :meth:`Tensor.sub`.
"""
return self.sub_(other, *args, **kwargs)
@doc_from_base()
@method_treelize()
def mul(self, other, *args, **kwargs):
"""
See :func:`treetensor.torch.mul`.
"""
return self.mul(other, *args, **kwargs)
@doc_from_base()
@return_self
@method_treelize()
def mul_(self, other, *args, **kwargs):
"""
In-place version of :meth:`Tensor.mul`.
"""
return self.mul_(other, *args, **kwargs)
@doc_from_base()
@method_treelize()
def div(self, other, *args, **kwargs):
"""
See :func:`treetensor.torch.div`.
"""
return self.div(other, *args, **kwargs)
@doc_from_base()
@return_self
@method_treelize()
def div_(self, other, *args, **kwargs):
"""
In-place version of :meth:`Tensor.div`.
"""
return self.div_(other, *args, **kwargs)
@doc_from_base()
@method_treelize()
def pow(self, exponent, *args, **kwargs):
"""
See :func:`treetensor.torch.pow`.
"""
return self.pow(exponent, *args, **kwargs)
@doc_from_base()
@return_self
@method_treelize()
def pow_(self, exponent, *args, **kwargs):
"""
In-place version of :meth:`Tensor.pow`.
"""
return self.pow_(exponent, *args, **kwargs)
@doc_from_base()
@method_treelize()
def neg(self, *args, **kwargs):
"""
See :func:`treetensor.torch.neg`.
"""
return self.neg(*args, **kwargs)
@doc_from_base()
@return_self
@method_treelize()
def neg_(self, *args, **kwargs):
"""
In-place version of :meth:`Tensor.neg`.
"""
return self.neg_(*args, **kwargs)
Markdown is supported
0% 或 .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑！
想要评论请 注册 或