Skip to content

P
pytensor
提交 66934167 authored 作者: Ricardo Vieira's avatar Ricardo Vieira 提交者： Ricardo Vieira
浏览文件
操作

Specialize TensorType string representation

上级 88516d2e
隐藏空白字符变更
内嵌 并排
正在显示 包含 68 行增加 和 62 行删除
pytensor/tensor/type.py
浏览文件 @ 66934167
......@@ -386,6 +386,8 @@ class TensorType(CType[np.ndarray], HasDataType, HasShape):
if self.name:
return self.name
else:
shape = self.shape
len_shape = len(shape)
def shape_str(s):
if s is None:
......@@ -393,14 +395,18 @@ class TensorType(CType[np.ndarray], HasDataType, HasShape):
else:
return str(s)
formatted_shape = ", ".join([shape_str(s) for s in self.shape])
if len(self.shape) == 1:
formatted_shape = ", ".join([shape_str(s) for s in shape])
if len_shape == 1:
formatted_shape += ","
return f"TensorType({self.dtype}, ({formatted_shape}))"
if len_shape > 2:
name = f"Tensor{len_shape}"
else:
name = ("Scalar", "Vector", "Matrix")[len_shape]
return f"{name}({self.dtype}, shape=({formatted_shape}))"
def __repr__(self):
return str(self)
return f"TensorType({self.dtype}, shape={self.shape})"
@staticmethod
def may_share_memory(a, b):
......
pytensor/tensor/var.py
浏览文件 @ 66934167
......@@ -1030,7 +1030,7 @@ class TensorConstant(TensorVariable, Constant[_TensorTypeType]):
else:
val = f"{self.data}"
if len(val) > 20:
val = val[:10] + ".." + val[-10:]
val = val[:10].strip() + " ... " + val[-10:].strip()
if self.name is not None:
name = self.name
......
tests/compile/test_builders.py
浏览文件 @ 66934167
......@@ -580,10 +580,10 @@ Inner graphs:
OpFromGraph{inline=False} [id A]
← Add [id E]
├─ *0-<TensorType(float64, (?, ?))> [id F]
├─ *0-<Matrix(float64, shape=(?, ?))> [id F]
└─ Mul [id G]
├─ *1-<TensorType(float64, (?, ?))> [id H]
└─ *2-<TensorType(float64, (?, ?))> [id I]
├─ *1-<Matrix(float64, shape=(?, ?))> [id H]
└─ *2-<Matrix(float64, shape=(?, ?))> [id I]
"""
for truth, out in zip(exp_res.split("\n"), lines):
......
tests/scan/test_printing.py
浏览文件 @ 66934167
......@@ -61,8 +61,8 @@ def test_debugprint_sitsot():
Scan{scan_fn, while_loop=False, inplace=none} [id C]
← Mul [id W] (inner_out_sit_sot-0)
├─ *0-<TensorType(float64, (?,))> [id X] -> [id E] (inner_in_sit_sot-0)
└─ *1-<TensorType(float64, (?,))> [id Y] -> [id M] (inner_in_non_seqs-0)"""
├─ *0-<Vector(float64, shape=(?,))> [id X] -> [id E] (inner_in_sit_sot-0)
└─ *1-<Vector(float64, shape=(?,))> [id Y] -> [id M] (inner_in_non_seqs-0)"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -119,8 +119,8 @@ def test_debugprint_sitsot_no_extra_info():
Scan{scan_fn, while_loop=False, inplace=none} [id C]
← Mul [id W]
├─ *0-<TensorType(float64, (?,))> [id X] -> [id E]
└─ *1-<TensorType(float64, (?,))> [id Y] -> [id M]"""
├─ *0-<Vector(float64, shape=(?,))> [id X] -> [id E]
└─ *1-<Vector(float64, shape=(?,))> [id Y] -> [id M]"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -185,10 +185,10 @@ def test_debugprint_nitsot():
Scan{scan_fn, while_loop=False, inplace=none} [id B]
← Mul [id X] (inner_out_nit_sot-0)
├─ *0-<TensorType(float64, ())> [id Y] -> [id S] (inner_in_seqs-0)
├─ *0-<Scalar(float64, shape=())> [id Y] -> [id S] (inner_in_seqs-0)
└─ Pow [id Z]
├─ *2-<TensorType(float64, ())> [id BA] -> [id W] (inner_in_non_seqs-0)
└─ *1-<TensorType(int64, ())> [id BB] -> [id U] (inner_in_seqs-1)"""
├─ *2-<Scalar(float64, shape=())> [id BA] -> [id W] (inner_in_non_seqs-0)
└─ *1-<Scalar(int64, shape=())> [id BB] -> [id U] (inner_in_seqs-1)"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -265,22 +265,22 @@ def test_debugprint_nested_scans():
Scan{scan_fn, while_loop=False, inplace=none} [id B]
← Mul [id Y] (inner_out_nit_sot-0)
├─ ExpandDims{axis=0} [id Z]
│ └─ *0-<TensorType(float64, ())> [id BA] -> [id S] (inner_in_seqs-0)
│ └─ *0-<Scalar(float64, shape=())> [id BA] -> [id S] (inner_in_seqs-0)
└─ Pow [id BB]
├─ Subtensor{i} [id BC]
│ ├─ Subtensor{start:} [id BD]
│ │ ├─ Scan{scan_fn, while_loop=False, inplace=none} [id BE] (outer_out_sit_sot-0)
│ │ │ ├─ *3-<TensorType(int32, ())> [id BF] -> [id X] (inner_in_non_seqs-1) (n_steps)
│ │ │ ├─ *3-<Scalar(int32, shape=())> [id BF] -> [id X] (inner_in_non_seqs-1) (n_steps)
│ │ │ ├─ SetSubtensor{:stop} [id BG] (outer_in_sit_sot-0)
│ │ │ │ ├─ AllocEmpty{dtype='float64'} [id BH]
│ │ │ │ │ ├─ Add [id BI]
│ │ │ │ │ │ ├─ *3-<TensorType(int32, ())> [id BF] -> [id X] (inner_in_non_seqs-1)
│ │ │ │ │ │ ├─ *3-<Scalar(int32, shape=())> [id BF] -> [id X] (inner_in_non_seqs-1)
│ │ │ │ │ │ └─ Subtensor{i} [id BJ]
│ │ │ │ │ │ ├─ Shape [id BK]
│ │ │ │ │ │ │ └─ Unbroadcast{0} [id BL]
│ │ │ │ │ │ │ └─ ExpandDims{axis=0} [id BM]
│ │ │ │ │ │ │ └─ Second [id BN]
│ │ │ │ │ │ │ ├─ *2-<TensorType(float64, (?,))> [id BO] -> [id W] (inner_in_non_seqs-0)
│ │ │ │ │ │ │ ├─ *2-<Vector(float64, shape=(?,))> [id BO] -> [id W] (inner_in_non_seqs-0)
│ │ │ │ │ │ │ └─ ExpandDims{axis=0} [id BP]
│ │ │ │ │ │ │ └─ TensorConstant{1.0} [id BQ]
│ │ │ │ │ │ └─ ScalarConstant{0} [id BR]
......@@ -294,16 +294,16 @@ def test_debugprint_nested_scans():
│ │ │ │ └─ ScalarFromTensor [id BV]
│ │ │ │ └─ Subtensor{i} [id BJ]
│ │ │ │ └─ ···
│ │ │ └─ *2-<TensorType(float64, (?,))> [id BO] -> [id W] (inner_in_non_seqs-0) (outer_in_non_seqs-0)
│ │ │ └─ *2-<Vector(float64, shape=(?,))> [id BO] -> [id W] (inner_in_non_seqs-0) (outer_in_non_seqs-0)
│ │ └─ ScalarConstant{1} [id BW]
│ └─ ScalarConstant{-1} [id BX]
└─ ExpandDims{axis=0} [id BY]
└─ *1-<TensorType(int64, ())> [id BZ] -> [id U] (inner_in_seqs-1)
└─ *1-<Scalar(int64, shape=())> [id BZ] -> [id U] (inner_in_seqs-1)
Scan{scan_fn, while_loop=False, inplace=none} [id BE]
← Mul [id CA] (inner_out_sit_sot-0)
├─ *0-<TensorType(float64, (?,))> [id CB] -> [id BG] (inner_in_sit_sot-0)
└─ *1-<TensorType(float64, (?,))> [id CC] -> [id BO] (inner_in_non_seqs-0)"""
├─ *0-<Vector(float64, shape=(?,))> [id CB] -> [id BG] (inner_in_sit_sot-0)
└─ *1-<Vector(float64, shape=(?,))> [id CC] -> [id BO] (inner_in_non_seqs-0)"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -356,28 +356,28 @@ def test_debugprint_nested_scans():
Inner graphs:
Scan{scan_fn, while_loop=False, inplace=none} [id E]
→ *0-<TensorType(float64, ())> [id Y] -> [id U] (inner_in_seqs-0)
→ *1-<TensorType(int64, ())> [id Z] -> [id W] (inner_in_seqs-1)
→ *2-<TensorType(float64, (?,))> [id BA] -> [id C] (inner_in_non_seqs-0)
→ *3-<TensorType(int32, ())> [id BB] -> [id B] (inner_in_non_seqs-1)
→ *0-<Scalar(float64, shape=())> [id Y] -> [id U] (inner_in_seqs-0)
→ *1-<Scalar(int64, shape=())> [id Z] -> [id W] (inner_in_seqs-1)
→ *2-<Vector(float64, shape=(?,))> [id BA] -> [id C] (inner_in_non_seqs-0)
→ *3-<Scalar(int32, shape=())> [id BB] -> [id B] (inner_in_non_seqs-1)
← Mul [id BC] (inner_out_nit_sot-0)
├─ ExpandDims{axis=0} [id BD]
│ └─ *0-<TensorType(float64, ())> [id Y] (inner_in_seqs-0)
│ └─ *0-<Scalar(float64, shape=())> [id Y] (inner_in_seqs-0)
└─ Pow [id BE]
├─ Subtensor{i} [id BF]
│ ├─ Subtensor{start:} [id BG]
│ │ ├─ Scan{scan_fn, while_loop=False, inplace=none} [id BH] (outer_out_sit_sot-0)
│ │ │ ├─ *3-<TensorType(int32, ())> [id BB] (inner_in_non_seqs-1) (n_steps)
│ │ │ ├─ *3-<Scalar(int32, shape=())> [id BB] (inner_in_non_seqs-1) (n_steps)
│ │ │ ├─ SetSubtensor{:stop} [id BI] (outer_in_sit_sot-0)
│ │ │ │ ├─ AllocEmpty{dtype='float64'} [id BJ]
│ │ │ │ │ ├─ Add [id BK]
│ │ │ │ │ │ ├─ *3-<TensorType(int32, ())> [id BB] (inner_in_non_seqs-1)
│ │ │ │ │ │ ├─ *3-<Scalar(int32, shape=())> [id BB] (inner_in_non_seqs-1)
│ │ │ │ │ │ └─ Subtensor{i} [id BL]
│ │ │ │ │ │ ├─ Shape [id BM]
│ │ │ │ │ │ │ └─ Unbroadcast{0} [id BN]
│ │ │ │ │ │ │ └─ ExpandDims{axis=0} [id BO]
│ │ │ │ │ │ │ └─ Second [id BP]
│ │ │ │ │ │ │ ├─ *2-<TensorType(float64, (?,))> [id BA] (inner_in_non_seqs-0)
│ │ │ │ │ │ │ ├─ *2-<Vector(float64, shape=(?,))> [id BA] (inner_in_non_seqs-0)
│ │ │ │ │ │ │ └─ ExpandDims{axis=0} [id BQ]
│ │ │ │ │ │ │ └─ TensorConstant{1.0} [id BR]
│ │ │ │ │ │ └─ ScalarConstant{0} [id BS]
......@@ -391,18 +391,18 @@ def test_debugprint_nested_scans():
│ │ │ │ └─ ScalarFromTensor [id BW]
│ │ │ │ └─ Subtensor{i} [id BL]
│ │ │ │ └─ ···
│ │ │ └─ *2-<TensorType(float64, (?,))> [id BA] (inner_in_non_seqs-0) (outer_in_non_seqs-0)
│ │ │ └─ *2-<Vector(float64, shape=(?,))> [id BA] (inner_in_non_seqs-0) (outer_in_non_seqs-0)
│ │ └─ ScalarConstant{1} [id BX]
│ └─ ScalarConstant{-1} [id BY]
└─ ExpandDims{axis=0} [id BZ]
└─ *1-<TensorType(int64, ())> [id Z] (inner_in_seqs-1)
└─ *1-<Scalar(int64, shape=())> [id Z] (inner_in_seqs-1)
Scan{scan_fn, while_loop=False, inplace=none} [id BH]
→ *0-<TensorType(float64, (?,))> [id CA] -> [id BI] (inner_in_sit_sot-0)
→ *1-<TensorType(float64, (?,))> [id CB] -> [id BA] (inner_in_non_seqs-0)
→ *0-<Vector(float64, shape=(?,))> [id CA] -> [id BI] (inner_in_sit_sot-0)
→ *1-<Vector(float64, shape=(?,))> [id CB] -> [id BA] (inner_in_non_seqs-0)
← Mul [id CC] (inner_out_sit_sot-0)
├─ *0-<TensorType(float64, (?,))> [id CA] (inner_in_sit_sot-0)
└─ *1-<TensorType(float64, (?,))> [id CB] (inner_in_non_seqs-0)"""
├─ *0-<Vector(float64, shape=(?,))> [id CA] (inner_in_sit_sot-0)
└─ *1-<Vector(float64, shape=(?,))> [id CB] (inner_in_non_seqs-0)"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -440,7 +440,7 @@ def test_debugprint_mitsot():
│ │ │ │ └─ Subtensor{i} [id H]
│ │ │ │ ├─ Shape [id I]
│ │ │ │ │ └─ Subtensor{:stop} [id J]
│ │ │ │ │ ├─ <TensorType(int64, (?,))> [id K]
│ │ │ │ │ ├─ <Vector(int64, shape=(?,))> [id K]
│ │ │ │ │ └─ ScalarConstant{2} [id L]
│ │ │ │ └─ ScalarConstant{0} [id M]
│ │ │ ├─ Subtensor{:stop} [id J]
......@@ -455,7 +455,7 @@ def test_debugprint_mitsot():
│ │ │ └─ Subtensor{i} [id R]
│ │ │ ├─ Shape [id S]
│ │ │ │ └─ Subtensor{:stop} [id T]
│ │ │ │ ├─ <TensorType(int64, (?,))> [id U]
│ │ │ │ ├─ <Vector(int64, shape=(?,))> [id U]
│ │ │ │ └─ ScalarConstant{2} [id V]
│ │ │ └─ ScalarConstant{0} [id W]
│ │ ├─ Subtensor{:stop} [id T]
......@@ -473,11 +473,11 @@ def test_debugprint_mitsot():
Scan{scan_fn, while_loop=False, inplace=none} [id C]
← Add [id BB] (inner_out_mit_sot-0)
├─ *1-<TensorType(int64, ())> [id BC] -> [id E] (inner_in_mit_sot-0-1)
└─ *0-<TensorType(int64, ())> [id BD] -> [id E] (inner_in_mit_sot-0-0)
├─ *1-<Scalar(int64, shape=())> [id BC] -> [id E] (inner_in_mit_sot-0-1)
└─ *0-<Scalar(int64, shape=())> [id BD] -> [id E] (inner_in_mit_sot-0-0)
← Add [id BE] (inner_out_mit_sot-1)
├─ *3-<TensorType(int64, ())> [id BF] -> [id O] (inner_in_mit_sot-1-1)
└─ *2-<TensorType(int64, ())> [id BG] -> [id O] (inner_in_mit_sot-1-0)"""
├─ *3-<Scalar(int64, shape=())> [id BF] -> [id O] (inner_in_mit_sot-1-1)
└─ *2-<Scalar(int64, shape=())> [id BG] -> [id O] (inner_in_mit_sot-1-0)"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -601,19 +601,19 @@ def test_debugprint_mitmot():
Scan{grad_of_scan_fn, while_loop=False, inplace=none} [id B]
← Add [id CM] (inner_out_mit_mot-0-0)
├─ Mul [id CN]
│ ├─ *2-<TensorType(float64, (?,))> [id CO] -> [id BL] (inner_in_mit_mot-0-0)
│ └─ *5-<TensorType(float64, (?,))> [id CP] -> [id P] (inner_in_non_seqs-0)
└─ *3-<TensorType(float64, (?,))> [id CQ] -> [id BL] (inner_in_mit_mot-0-1)
│ ├─ *2-<Vector(float64, shape=(?,))> [id CO] -> [id BL] (inner_in_mit_mot-0-0)
│ └─ *5-<Vector(float64, shape=(?,))> [id CP] -> [id P] (inner_in_non_seqs-0)
└─ *3-<Vector(float64, shape=(?,))> [id CQ] -> [id BL] (inner_in_mit_mot-0-1)
← Add [id CR] (inner_out_sit_sot-0)
├─ Mul [id CS]
│ ├─ *2-<TensorType(float64, (?,))> [id CO] -> [id BL] (inner_in_mit_mot-0-0)
│ └─ *0-<TensorType(float64, (?,))> [id CT] -> [id Z] (inner_in_seqs-0)
└─ *4-<TensorType(float64, (?,))> [id CU] -> [id CE] (inner_in_sit_sot-0)
│ ├─ *2-<Vector(float64, shape=(?,))> [id CO] -> [id BL] (inner_in_mit_mot-0-0)
│ └─ *0-<Vector(float64, shape=(?,))> [id CT] -> [id Z] (inner_in_seqs-0)
└─ *4-<Vector(float64, shape=(?,))> [id CU] -> [id CE] (inner_in_sit_sot-0)
Scan{scan_fn, while_loop=False, inplace=none} [id F]
← Mul [id CV] (inner_out_sit_sot-0)
├─ *0-<TensorType(float64, (?,))> [id CT] -> [id H] (inner_in_sit_sot-0)
└─ *1-<TensorType(float64, (?,))> [id CW] -> [id P] (inner_in_non_seqs-0)"""
├─ *0-<Vector(float64, shape=(?,))> [id CT] -> [id H] (inner_in_sit_sot-0)
└─ *1-<Vector(float64, shape=(?,))> [id CW] -> [id P] (inner_in_non_seqs-0)"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
......@@ -643,13 +643,13 @@ def test_debugprint_compiled_fn():
expected_output = """Scan{scan_fn, while_loop=False, inplace=all} [id A] 2 (outer_out_sit_sot-0)
├─ TensorConstant{20000} [id B] (n_steps)
├─ TensorConstant{[ 0 ..998 19999]} [id C] (outer_in_seqs-0)
├─ TensorConstant{[ 0 ... 998 19999]} [id C] (outer_in_seqs-0)
├─ SetSubtensor{:stop} [id D] 1 (outer_in_sit_sot-0)
│ ├─ AllocEmpty{dtype='int64'} [id E] 0
│ │ └─ TensorConstant{20000} [id B]
│ ├─ TensorConstant{(1,) of 0} [id F]
│ └─ ScalarConstant{1} [id G]
└─ <TensorType(float64, (20000, 2, 2))> [id H] (outer_in_non_seqs-0)
└─ <Tensor3(float64, shape=(20000, 2, 2))> [id H] (outer_in_non_seqs-0)
Inner graphs:
......@@ -657,11 +657,11 @@ def test_debugprint_compiled_fn():
← Composite{switch(lt(i0, i1), i2, i0)} [id I] (inner_out_sit_sot-0)
├─ TensorConstant{0} [id J]
├─ Subtensor{i, j, k} [id K]
│ ├─ *2-<TensorType(float64, (20000, 2, 2))> [id L] -> [id H] (inner_in_non_seqs-0)
│ ├─ *2-<Tensor3(float64, shape=(20000, 2, 2))> [id L] -> [id H] (inner_in_non_seqs-0)
│ ├─ ScalarFromTensor [id M]
│ │ └─ *0-<TensorType(int64, ())> [id N] -> [id C] (inner_in_seqs-0)
│ │ └─ *0-<Scalar(int64, shape=())> [id N] -> [id C] (inner_in_seqs-0)
│ ├─ ScalarFromTensor [id O]
│ │ └─ *1-<TensorType(int64, ())> [id P] -> [id D] (inner_in_sit_sot-0)
│ │ └─ *1-<Scalar(int64, shape=())> [id P] -> [id D] (inner_in_sit_sot-0)
│ └─ ScalarConstant{0} [id Q]
└─ TensorConstant{1} [id R]
......
tests/tensor/test_type.py
浏览文件 @ 66934167
......@@ -252,7 +252,7 @@ def test_fixed_shape_basic():
assert t1.shape == (2, 3)
assert t1.broadcastable == (False, False)
assert str(t1) == "TensorType(float64, (2, 3))"
assert str(t1) == "Matrix(float64, shape=(2, 3))"
t1 = TensorType("float64", shape=(1,))
assert t1.shape == (1,)
......
tests/test_printing.py
浏览文件 @ 66934167
......@@ -282,7 +282,7 @@ def test_debugprint():
│ │ └─ B
│ ├─ TensorConstant{1.0}
│ ├─ B
│ ├─ <TensorType(float64, (?,))>
│ ├─ <Vector(float64, shape=(?,))>
│ └─ TensorConstant{0.0}
├─ D
└─ A
......@@ -316,9 +316,9 @@ def test_debugprint_id_type():
exp_res = f"""Add [id {e_at.auto_name}]
├─ dot [id {d_at.auto_name}]
│ ├─ <TensorType(float64, (?, ?))> [id {b_at.auto_name}]
│ └─ <TensorType(float64, (?,))> [id {a_at.auto_name}]
└─ <TensorType(float64, (?,))> [id {a_at.auto_name}]
│ ├─ <Matrix(float64, shape=(?, ?))> [id {b_at.auto_name}]
│ └─ <Vector(float64, shape=(?,))> [id {a_at.auto_name}]
└─ <Vector(float64, shape=(?,))> [id {a_at.auto_name}]
"""
assert [l.strip() for l in s.split("\n")] == [
......@@ -329,7 +329,7 @@ def test_debugprint_id_type():
def test_pprint():
x = dvector()
y = x[1]
assert pp(y) == "<TensorType(float64, (?,))>[1]"
assert pp(y) == "<Vector(float64, shape=(?,))>[1]"
def test_debugprint_inner_graph():
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论