Skip to content

P
pytensor
提交 89b0b822 authored 作者: Brandon T. Willard's avatar Brandon T. Willard 提交者： Brandon T. Willard
浏览文件
操作

Move Scan printing tests into tests.scan.test_printing

上级 93502185
隐藏空白字符变更
内嵌 并排
正在显示 包含 465 行增加 和 459 行删除
tests/scan/test_printing.py 0 → 100644
浏览文件 @ 89b0b822
import numpy as np
import pytest
import aesara
import aesara.tensor as aet
from aesara.printing import debugprint, pydot_imported, pydotprint
from aesara.tensor.type import dvector, iscalar, scalar, vector
def test_scan_debugprint1():
k = iscalar("k")
A = dvector("A")
# Symbolic description of the result
result, updates = aesara.scan(
fn=lambda prior_result, A: prior_result * A,
outputs_info=aet.ones_like(A),
non_sequences=A,
n_steps=k,
)
final_result = result[-1]
output_str = debugprint(final_result, file="str")
lines = output_str.split("\n")
expected_output = """Subtensor{int64} [id A] ''
|Subtensor{int64::} [id B] ''
| |for{cpu,scan_fn} [id C] ''
| | |k [id D]
| | |IncSubtensor{Set;:int64:} [id E] ''
| | | |AllocEmpty{dtype='float64'} [id F] ''
| | | | |Elemwise{add,no_inplace} [id G] ''
| | | | | |k [id D]
| | | | | |Subtensor{int64} [id H] ''
| | | | | |Shape [id I] ''
| | | | | | |Rebroadcast{0} [id J] ''
| | | | | | |InplaceDimShuffle{x,0} [id K] ''
| | | | | | |Elemwise{second,no_inplace} [id L] ''
| | | | | | |A [id M]
| | | | | | |InplaceDimShuffle{x} [id N] ''
| | | | | | |TensorConstant{1.0} [id O]
| | | | | |ScalarConstant{0} [id P]
| | | | |Subtensor{int64} [id Q] ''
| | | | |Shape [id R] ''
| | | | | |Rebroadcast{0} [id J] ''
| | | | |ScalarConstant{1} [id S]
| | | |Rebroadcast{0} [id J] ''
| | | |ScalarFromTensor [id T] ''
| | | |Subtensor{int64} [id H] ''
| | |A [id M]
| |ScalarConstant{1} [id U]
|ScalarConstant{-1} [id V]
Inner graphs of the scan ops:
for{cpu,scan_fn} [id C] ''
>Elemwise{mul,no_inplace} [id W] ''
> |<TensorType(float64, vector)> [id X] -> [id E]
> |A_copy [id Y] -> [id M]"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
def test_scan_debugprint2():
coefficients = vector("coefficients")
x = scalar("x")
max_coefficients_supported = 10000
# Generate the components of the polynomial
components, updates = aesara.scan(
fn=lambda coefficient, power, free_variable: coefficient
* (free_variable ** power),
outputs_info=None,
sequences=[coefficients, aet.arange(max_coefficients_supported)],
non_sequences=x,
)
# Sum them up
polynomial = components.sum()
output_str = debugprint(polynomial, file="str")
lines = output_str.split("\n")
expected_output = """Sum{acc_dtype=float64} [id A] ''
|for{cpu,scan_fn} [id B] ''
|Elemwise{scalar_minimum,no_inplace} [id C] ''
| |Subtensor{int64} [id D] ''
| | |Shape [id E] ''
| | | |Subtensor{int64::} [id F] 'coefficients[0:]'
| | | |coefficients [id G]
| | | |ScalarConstant{0} [id H]
| | |ScalarConstant{0} [id I]
| |Subtensor{int64} [id J] ''
| |Shape [id K] ''
| | |Subtensor{int64::} [id L] ''
| | |ARange{dtype='int64'} [id M] ''
| | | |TensorConstant{0} [id N]
| | | |TensorConstant{10000} [id O]
| | | |TensorConstant{1} [id P]
| | |ScalarConstant{0} [id Q]
| |ScalarConstant{0} [id R]
|Subtensor{:int64:} [id S] ''
| |Subtensor{int64::} [id F] 'coefficients[0:]'
| |ScalarFromTensor [id T] ''
| |Elemwise{scalar_minimum,no_inplace} [id C] ''
|Subtensor{:int64:} [id U] ''
| |Subtensor{int64::} [id L] ''
| |ScalarFromTensor [id V] ''
| |Elemwise{scalar_minimum,no_inplace} [id C] ''
|Elemwise{scalar_minimum,no_inplace} [id C] ''
|x [id W]
Inner graphs of the scan ops:
for{cpu,scan_fn} [id B] ''
>Elemwise{mul,no_inplace} [id X] ''
> |coefficients[t] [id Y] -> [id S]
> |Elemwise{pow,no_inplace} [id Z] ''
> |x_copy [id BA] -> [id W]
> |<TensorType(int64, scalar)> [id BB] -> [id U]"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
@aesara.config.change_flags(optimizer_verbose=True)
def test_scan_debugprint3():
coefficients = dvector("coefficients")
max_coefficients_supported = 10
k = iscalar("k")
A = dvector("A")
# compute A**k
def compute_A_k(A, k):
# Symbolic description of the result
result, updates = aesara.scan(
fn=lambda prior_result, A: prior_result * A,
outputs_info=aet.ones_like(A),
non_sequences=A,
n_steps=k,
)
A_k = result[-1]
return A_k
# Generate the components of the polynomial
components, updates = aesara.scan(
fn=lambda coefficient, power, some_A, some_k: coefficient
* (compute_A_k(some_A, some_k) ** power),
outputs_info=None,
sequences=[coefficients, aet.arange(max_coefficients_supported)],
non_sequences=[A, k],
)
# Sum them up
polynomial = components.sum()
final_result = polynomial
output_str = debugprint(final_result, file="str")
lines = output_str.split("\n")
expected_output = """Sum{acc_dtype=float64} [id A] ''
|for{cpu,scan_fn} [id B] ''
|Elemwise{scalar_minimum,no_inplace} [id C] ''
| |Subtensor{int64} [id D] ''
| | |Shape [id E] ''
| | | |Subtensor{int64::} [id F] 'coefficients[0:]'
| | | |coefficients [id G]
| | | |ScalarConstant{0} [id H]
| | |ScalarConstant{0} [id I]
| |Subtensor{int64} [id J] ''
| |Shape [id K] ''
| | |Subtensor{int64::} [id L] ''
| | |ARange{dtype='int64'} [id M] ''
| | | |TensorConstant{0} [id N]
| | | |TensorConstant{10} [id O]
| | | |TensorConstant{1} [id P]
| | |ScalarConstant{0} [id Q]
| |ScalarConstant{0} [id R]
|Subtensor{:int64:} [id S] ''
| |Subtensor{int64::} [id F] 'coefficients[0:]'
| |ScalarFromTensor [id T] ''
| |Elemwise{scalar_minimum,no_inplace} [id C] ''
|Subtensor{:int64:} [id U] ''
| |Subtensor{int64::} [id L] ''
| |ScalarFromTensor [id V] ''
| |Elemwise{scalar_minimum,no_inplace} [id C] ''
|Elemwise{scalar_minimum,no_inplace} [id C] ''
|A [id W]
|k [id X]
Inner graphs of the scan ops:
for{cpu,scan_fn} [id B] ''
>Elemwise{mul,no_inplace} [id Y] ''
> |InplaceDimShuffle{x} [id Z] ''
> | |coefficients[t] [id BA] -> [id S]
> |Elemwise{pow,no_inplace} [id BB] ''
> |Subtensor{int64} [id BC] ''
> | |Subtensor{int64::} [id BD] ''
> | | |for{cpu,scan_fn} [id BE] ''
> | | | |k_copy [id BF] -> [id X]
> | | | |IncSubtensor{Set;:int64:} [id BG] ''
> | | | | |AllocEmpty{dtype='float64'} [id BH] ''
> | | | | | |Elemwise{add,no_inplace} [id BI] ''
> | | | | | | |k_copy [id BF] -> [id X]
> | | | | | | |Subtensor{int64} [id BJ] ''
> | | | | | | |Shape [id BK] ''
> | | | | | | | |Rebroadcast{0} [id BL] ''
> | | | | | | | |InplaceDimShuffle{x,0} [id BM] ''
> | | | | | | | |Elemwise{second,no_inplace} [id BN] ''
> | | | | | | | |A_copy [id BO] -> [id W]
> | | | | | | | |InplaceDimShuffle{x} [id BP] ''
> | | | | | | | |TensorConstant{1.0} [id BQ]
> | | | | | | |ScalarConstant{0} [id BR]
> | | | | | |Subtensor{int64} [id BS] ''
> | | | | | |Shape [id BT] ''
> | | | | | | |Rebroadcast{0} [id BL] ''
> | | | | | |ScalarConstant{1} [id BU]
> | | | | |Rebroadcast{0} [id BL] ''
> | | | | |ScalarFromTensor [id BV] ''
> | | | | |Subtensor{int64} [id BJ] ''
> | | | |A_copy [id BO] -> [id W]
> | | |ScalarConstant{1} [id BW]
> | |ScalarConstant{-1} [id BX]
> |InplaceDimShuffle{x} [id BY] ''
> |<TensorType(int64, scalar)> [id BZ] -> [id U]
for{cpu,scan_fn} [id BE] ''
>Elemwise{mul,no_inplace} [id CA] ''
> |<TensorType(float64, vector)> [id CB] -> [id BG]
> |A_copy [id CC] -> [id BO]"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
def test_scan_debugprint4():
def fn(a_m2, a_m1, b_m2, b_m1):
return a_m1 + a_m2, b_m1 + b_m2
a0 = aesara.shared(np.arange(2, dtype="int64"))
b0 = aesara.shared(np.arange(2, dtype="int64"))
(a, b), _ = aesara.scan(
fn,
outputs_info=[
{"initial": a0, "taps": [-2, -1]},
{"initial": b0, "taps": [-2, -1]},
],
n_steps=5,
)
final_result = a + b
output_str = debugprint(final_result, file="str")
lines = output_str.split("\n")
expected_output = """Elemwise{add,no_inplace} [id A] ''
|Subtensor{int64::} [id B] ''
| |for{cpu,scan_fn}.0 [id C] ''
| | |TensorConstant{5} [id D]
| | |IncSubtensor{Set;:int64:} [id E] ''
| | | |AllocEmpty{dtype='int64'} [id F] ''
| | | | |Elemwise{add,no_inplace} [id G] ''
| | | | |TensorConstant{5} [id D]
| | | | |Subtensor{int64} [id H] ''
| | | | |Shape [id I] ''
| | | | | |Subtensor{:int64:} [id J] ''
| | | | | |<TensorType(int64, vector)> [id K]
| | | | | |ScalarConstant{2} [id L]
| | | | |ScalarConstant{0} [id M]
| | | |Subtensor{:int64:} [id J] ''
| | | |ScalarFromTensor [id N] ''
| | | |Subtensor{int64} [id H] ''
| | |IncSubtensor{Set;:int64:} [id O] ''
| | |AllocEmpty{dtype='int64'} [id P] ''
| | | |Elemwise{add,no_inplace} [id Q] ''
| | | |TensorConstant{5} [id D]
| | | |Subtensor{int64} [id R] ''
| | | |Shape [id S] ''
| | | | |Subtensor{:int64:} [id T] ''
| | | | |<TensorType(int64, vector)> [id U]
| | | | |ScalarConstant{2} [id V]
| | | |ScalarConstant{0} [id W]
| | |Subtensor{:int64:} [id T] ''
| | |ScalarFromTensor [id X] ''
| | |Subtensor{int64} [id R] ''
| |ScalarConstant{2} [id Y]
|Subtensor{int64::} [id Z] ''
|for{cpu,scan_fn}.1 [id C] ''
|ScalarConstant{2} [id BA]
Inner graphs of the scan ops:
for{cpu,scan_fn}.0 [id C] ''
>Elemwise{add,no_inplace} [id BB] ''
> |<TensorType(int64, scalar)> [id BC] -> [id E]
> |<TensorType(int64, scalar)> [id BD] -> [id E]
>Elemwise{add,no_inplace} [id BE] ''
> |<TensorType(int64, scalar)> [id BF] -> [id O]
> |<TensorType(int64, scalar)> [id BG] -> [id O]
for{cpu,scan_fn}.1 [id C] ''
>Elemwise{add,no_inplace} [id BB] ''
>Elemwise{add,no_inplace} [id BE] ''"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
def test_scan_debugprint5():
k = iscalar("k")
A = dvector("A")
# Symbolic description of the result
result, updates = aesara.scan(
fn=lambda prior_result, A: prior_result * A,
outputs_info=aet.ones_like(A),
non_sequences=A,
n_steps=k,
)
final_result = aesara.grad(result[-1].sum(), A)
output_str = debugprint(final_result, file="str")
lines = output_str.split("\n")
expected_output = """Subtensor{int64} [id A] ''
|for{cpu,grad_of_scan_fn}.1 [id B] ''
| |Elemwise{sub,no_inplace} [id C] ''
| | |Subtensor{int64} [id D] ''
| | | |Shape [id E] ''
| | | | |for{cpu,scan_fn} [id F] ''
| | | | |k [id G]
| | | | |IncSubtensor{Set;:int64:} [id H] ''
| | | | | |AllocEmpty{dtype='float64'} [id I] ''
| | | | | | |Elemwise{add,no_inplace} [id J] ''
| | | | | | | |k [id G]
| | | | | | | |Subtensor{int64} [id K] ''
| | | | | | | |Shape [id L] ''
| | | | | | | | |Rebroadcast{0} [id M] ''
| | | | | | | | |InplaceDimShuffle{x,0} [id N] ''
| | | | | | | | |Elemwise{second,no_inplace} [id O] ''
| | | | | | | | |A [id P]
| | | | | | | | |InplaceDimShuffle{x} [id Q] ''
| | | | | | | | |TensorConstant{1.0} [id R]
| | | | | | | |ScalarConstant{0} [id S]
| | | | | | |Subtensor{int64} [id T] ''
| | | | | | |Shape [id U] ''
| | | | | | | |Rebroadcast{0} [id M] ''
| | | | | | |ScalarConstant{1} [id V]
| | | | | |Rebroadcast{0} [id M] ''
| | | | | |ScalarFromTensor [id W] ''
| | | | | |Subtensor{int64} [id K] ''
| | | | |A [id P]
| | | |ScalarConstant{0} [id X]
| | |TensorConstant{1} [id Y]
| |Subtensor{:int64:} [id Z] ''
| | |Subtensor{::int64} [id BA] ''
| | | |Subtensor{:int64:} [id BB] ''
| | | | |for{cpu,scan_fn} [id F] ''
| | | | |ScalarConstant{-1} [id BC]
| | | |ScalarConstant{-1} [id BD]
| | |ScalarFromTensor [id BE] ''
| | |Elemwise{sub,no_inplace} [id C] ''
| |Subtensor{:int64:} [id BF] ''
| | |Subtensor{:int64:} [id BG] ''
| | | |Subtensor{::int64} [id BH] ''
| | | | |for{cpu,scan_fn} [id F] ''
| | | | |ScalarConstant{-1} [id BI]
| | | |ScalarConstant{-1} [id BJ]
| | |ScalarFromTensor [id BK] ''
| | |Elemwise{sub,no_inplace} [id C] ''
| |Subtensor{::int64} [id BL] ''
| | |IncSubtensor{Inc;int64::} [id BM] ''
| | | |Elemwise{second,no_inplace} [id BN] ''
| | | | |for{cpu,scan_fn} [id F] ''
| | | | |InplaceDimShuffle{x,x} [id BO] ''
| | | | |TensorConstant{0.0} [id BP]
| | | |IncSubtensor{Inc;int64} [id BQ] ''
| | | | |Elemwise{second,no_inplace} [id BR] ''
| | | | | |Subtensor{int64::} [id BS] ''
| | | | | | |for{cpu,scan_fn} [id F] ''
| | | | | | |ScalarConstant{1} [id BT]
| | | | | |InplaceDimShuffle{x,x} [id BU] ''
| | | | | |TensorConstant{0.0} [id BV]
| | | | |Elemwise{second} [id BW] ''
| | | | | |Subtensor{int64} [id BX] ''
| | | | | | |Subtensor{int64::} [id BS] ''
| | | | | | |ScalarConstant{-1} [id BY]
| | | | | |InplaceDimShuffle{x} [id BZ] ''
| | | | | |Elemwise{second,no_inplace} [id CA] ''
| | | | | |Sum{acc_dtype=float64} [id CB] ''
| | | | | | |Subtensor{int64} [id BX] ''
| | | | | |TensorConstant{1.0} [id CC]
| | | | |ScalarConstant{-1} [id BY]
| | | |ScalarConstant{1} [id BT]
| | |ScalarConstant{-1} [id CD]
| |Alloc [id CE] ''
| | |TensorConstant{0.0} [id CF]
| | |Elemwise{add,no_inplace} [id CG] ''
| | | |Elemwise{sub,no_inplace} [id C] ''
| | | |TensorConstant{1} [id CH]
| | |Subtensor{int64} [id CI] ''
| | |Shape [id CJ] ''
| | | |A [id P]
| | |ScalarConstant{0} [id CK]
| |A [id P]
|ScalarConstant{-1} [id CL]
Inner graphs of the scan ops:
for{cpu,grad_of_scan_fn}.1 [id B] ''
>Elemwise{add,no_inplace} [id CM] ''
> |Elemwise{mul} [id CN] ''
> | |<TensorType(float64, vector)> [id CO] -> [id BL]
> | |A_copy [id CP] -> [id P]
> |<TensorType(float64, vector)> [id CQ] -> [id BL]
>Elemwise{add,no_inplace} [id CR] ''
> |Elemwise{mul} [id CS] ''
> | |<TensorType(float64, vector)> [id CO] -> [id BL]
> | |<TensorType(float64, vector)> [id CT] -> [id Z]
> |<TensorType(float64, vector)> [id CU] -> [id CE]
for{cpu,scan_fn} [id F] ''
>Elemwise{mul,no_inplace} [id CV] ''
> |<TensorType(float64, vector)> [id CT] -> [id H]
> |A_copy [id CP] -> [id P]
for{cpu,scan_fn} [id F] ''
>Elemwise{mul,no_inplace} [id CV] ''
for{cpu,scan_fn} [id F] ''
>Elemwise{mul,no_inplace} [id CV] ''
for{cpu,scan_fn} [id F] ''
>Elemwise{mul,no_inplace} [id CV] ''
for{cpu,scan_fn} [id F] ''
>Elemwise{mul,no_inplace} [id CV] ''"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
@pytest.mark.skipif(not pydot_imported, reason="pydot not available")
def test_printing_scan():
def f_pow2(x_tm1):
return 2 * x_tm1
state = scalar("state")
n_steps = iscalar("nsteps")
output, updates = aesara.scan(
f_pow2, [], state, [], n_steps=n_steps, truncate_gradient=-1, go_backwards=False
)
f = aesara.function(
[state, n_steps], output, updates=updates, allow_input_downcast=True
)
pydotprint(output, scan_graphs=True)
pydotprint(f, scan_graphs=True)
tests/test_printing.py
浏览文件 @ 89b0b822
......@@ -4,11 +4,9 @@ Tests of printing functionality
import logging
from io import StringIO
import numpy as np
import pytest
import aesara
import aesara.tensor as aet
from aesara.printing import (
debugprint,
min_informative_str,
......@@ -16,7 +14,7 @@ from aesara.printing import (
pydot_imported,
pydotprint,
)
from aesara.tensor.type import dvector, iscalar, matrix, scalar, vector
from aesara.tensor.type import dvector, matrix
@pytest.mark.skipif(not pydot_imported, reason="pydot not available")
......@@ -252,462 +250,6 @@ def test_debugprint():
assert s == reference
def test_scan_debugprint1():
k = iscalar("k")
A = dvector("A")
# Symbolic description of the result
result, updates = aesara.scan(
fn=lambda prior_result, A: prior_result * A,
outputs_info=aet.ones_like(A),
non_sequences=A,
n_steps=k,
)
final_result = result[-1]
output_str = debugprint(final_result, file="str")
lines = output_str.split("\n")
expected_output = """Subtensor{int64} [id A] ''
|Subtensor{int64::} [id B] ''
| |for{cpu,scan_fn} [id C] ''
| | |k [id D]
| | |IncSubtensor{Set;:int64:} [id E] ''
| | | |AllocEmpty{dtype='float64'} [id F] ''
| | | | |Elemwise{add,no_inplace} [id G] ''
| | | | | |k [id D]
| | | | | |Subtensor{int64} [id H] ''
| | | | | |Shape [id I] ''
| | | | | | |Rebroadcast{0} [id J] ''
| | | | | | |InplaceDimShuffle{x,0} [id K] ''
| | | | | | |Elemwise{second,no_inplace} [id L] ''
| | | | | | |A [id M]
| | | | | | |InplaceDimShuffle{x} [id N] ''
| | | | | | |TensorConstant{1.0} [id O]
| | | | | |ScalarConstant{0} [id P]
| | | | |Subtensor{int64} [id Q] ''
| | | | |Shape [id R] ''
| | | | | |Rebroadcast{0} [id J] ''
| | | | |ScalarConstant{1} [id S]
| | | |Rebroadcast{0} [id J] ''
| | | |ScalarFromTensor [id T] ''
| | | |Subtensor{int64} [id H] ''
| | |A [id M]
| |ScalarConstant{1} [id U]
|ScalarConstant{-1} [id V]
Inner graphs of the scan ops:
for{cpu,scan_fn} [id C] ''
>Elemwise{mul,no_inplace} [id W] ''
> |<TensorType(float64, vector)> [id X] -> [id E]
> |A_copy [id Y] -> [id M]"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
def test_scan_debugprint2():
coefficients = vector("coefficients")
x = scalar("x")
max_coefficients_supported = 10000
# Generate the components of the polynomial
components, updates = aesara.scan(
fn=lambda coefficient, power, free_variable: coefficient
* (free_variable ** power),
outputs_info=None,
sequences=[coefficients, aet.arange(max_coefficients_supported)],
non_sequences=x,
)
# Sum them up
polynomial = components.sum()
output_str = debugprint(polynomial, file="str")
lines = output_str.split("\n")
expected_output = """Sum{acc_dtype=float64} [id A] ''
|for{cpu,scan_fn} [id B] ''
|Elemwise{scalar_minimum,no_inplace} [id C] ''
| |Subtensor{int64} [id D] ''
| | |Shape [id E] ''
| | | |Subtensor{int64::} [id F] 'coefficients[0:]'
| | | |coefficients [id G]
| | | |ScalarConstant{0} [id H]
| | |ScalarConstant{0} [id I]
| |Subtensor{int64} [id J] ''
| |Shape [id K] ''
| | |Subtensor{int64::} [id L] ''
| | |ARange{dtype='int64'} [id M] ''
| | | |TensorConstant{0} [id N]
| | | |TensorConstant{10000} [id O]
| | | |TensorConstant{1} [id P]
| | |ScalarConstant{0} [id Q]
| |ScalarConstant{0} [id R]
|Subtensor{:int64:} [id S] ''
| |Subtensor{int64::} [id F] 'coefficients[0:]'
| |ScalarFromTensor [id T] ''
| |Elemwise{scalar_minimum,no_inplace} [id C] ''
|Subtensor{:int64:} [id U] ''
| |Subtensor{int64::} [id L] ''
| |ScalarFromTensor [id V] ''
| |Elemwise{scalar_minimum,no_inplace} [id C] ''
|Elemwise{scalar_minimum,no_inplace} [id C] ''
|x [id W]
Inner graphs of the scan ops:
for{cpu,scan_fn} [id B] ''
>Elemwise{mul,no_inplace} [id X] ''
> |coefficients[t] [id Y] -> [id S]
> |Elemwise{pow,no_inplace} [id Z] ''
> |x_copy [id BA] -> [id W]
> |<TensorType(int64, scalar)> [id BB] -> [id U]"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
def test_scan_debugprint3():
coefficients = dvector("coefficients")
max_coefficients_supported = 10
k = iscalar("k")
A = dvector("A")
# compute A**k
def compute_A_k(A, k):
# Symbolic description of the result
result, updates = aesara.scan(
fn=lambda prior_result, A: prior_result * A,
outputs_info=aet.ones_like(A),
non_sequences=A,
n_steps=k,
)
A_k = result[-1]
return A_k
# Generate the components of the polynomial
components, updates = aesara.scan(
fn=lambda coefficient, power, some_A, some_k: coefficient
* (compute_A_k(some_A, some_k) ** power),
outputs_info=None,
sequences=[coefficients, aet.arange(max_coefficients_supported)],
non_sequences=[A, k],
)
# Sum them up
polynomial = components.sum()
final_result = polynomial
output_str = debugprint(final_result, file="str")
lines = output_str.split("\n")
expected_output = """Sum{acc_dtype=float64} [id A] ''
|for{cpu,scan_fn} [id B] ''
|Elemwise{scalar_minimum,no_inplace} [id C] ''
| |Subtensor{int64} [id D] ''
| | |Shape [id E] ''
| | | |Subtensor{int64::} [id F] 'coefficients[0:]'
| | | |coefficients [id G]
| | | |ScalarConstant{0} [id H]
| | |ScalarConstant{0} [id I]
| |Subtensor{int64} [id J] ''
| |Shape [id K] ''
| | |Subtensor{int64::} [id L] ''
| | |ARange{dtype='int64'} [id M] ''
| | | |TensorConstant{0} [id N]
| | | |TensorConstant{10} [id O]
| | | |TensorConstant{1} [id P]
| | |ScalarConstant{0} [id Q]
| |ScalarConstant{0} [id R]
|Subtensor{:int64:} [id S] ''
| |Subtensor{int64::} [id F] 'coefficients[0:]'
| |ScalarFromTensor [id T] ''
| |Elemwise{scalar_minimum,no_inplace} [id C] ''
|Subtensor{:int64:} [id U] ''
| |Subtensor{int64::} [id L] ''
| |ScalarFromTensor [id V] ''
| |Elemwise{scalar_minimum,no_inplace} [id C] ''
|Elemwise{scalar_minimum,no_inplace} [id C] ''
|A [id W]
|k [id X]
Inner graphs of the scan ops:
for{cpu,scan_fn} [id B] ''
>Elemwise{mul,no_inplace} [id Y] ''
> |InplaceDimShuffle{x} [id Z] ''
> | |coefficients[t] [id BA] -> [id S]
> |Elemwise{pow,no_inplace} [id BB] ''
> |Subtensor{int64} [id BC] ''
> | |Subtensor{int64::} [id BD] ''
> | | |for{cpu,scan_fn} [id BE] ''
> | | | |k_copy [id BF] -> [id X]
> | | | |IncSubtensor{Set;:int64:} [id BG] ''
> | | | | |AllocEmpty{dtype='float64'} [id BH] ''
> | | | | | |Elemwise{add,no_inplace} [id BI] ''
> | | | | | | |k_copy [id BF] -> [id X]
> | | | | | | |Subtensor{int64} [id BJ] ''
> | | | | | | |Shape [id BK] ''
> | | | | | | | |Rebroadcast{0} [id BL] ''
> | | | | | | | |InplaceDimShuffle{x,0} [id BM] ''
> | | | | | | | |Elemwise{second,no_inplace} [id BN] ''
> | | | | | | | |A_copy [id BO] -> [id W]
> | | | | | | | |InplaceDimShuffle{x} [id BP] ''
> | | | | | | | |TensorConstant{1.0} [id BQ]
> | | | | | | |ScalarConstant{0} [id BR]
> | | | | | |Subtensor{int64} [id BS] ''
> | | | | | |Shape [id BT] ''
> | | | | | | |Rebroadcast{0} [id BL] ''
> | | | | | |ScalarConstant{1} [id BU]
> | | | | |Rebroadcast{0} [id BL] ''
> | | | | |ScalarFromTensor [id BV] ''
> | | | | |Subtensor{int64} [id BJ] ''
> | | | |A_copy [id BO] -> [id W]
> | | |ScalarConstant{1} [id BW]
> | |ScalarConstant{-1} [id BX]
> |InplaceDimShuffle{x} [id BY] ''
> |<TensorType(int64, scalar)> [id BZ] -> [id U]
for{cpu,scan_fn} [id BE] ''
>Elemwise{mul,no_inplace} [id CA] ''
> |<TensorType(float64, vector)> [id CB] -> [id BG]
> |A_copy [id CC] -> [id BO]"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
def test_scan_debugprint4():
def fn(a_m2, a_m1, b_m2, b_m1):
return a_m1 + a_m2, b_m1 + b_m2
a0 = aesara.shared(np.arange(2, dtype="int64"))
b0 = aesara.shared(np.arange(2, dtype="int64"))
(a, b), _ = aesara.scan(
fn,
outputs_info=[
{"initial": a0, "taps": [-2, -1]},
{"initial": b0, "taps": [-2, -1]},
],
n_steps=5,
)
final_result = a + b
output_str = debugprint(final_result, file="str")
lines = output_str.split("\n")
expected_output = """Elemwise{add,no_inplace} [id A] ''
|Subtensor{int64::} [id B] ''
| |for{cpu,scan_fn}.0 [id C] ''
| | |TensorConstant{5} [id D]
| | |IncSubtensor{Set;:int64:} [id E] ''
| | | |AllocEmpty{dtype='int64'} [id F] ''
| | | | |Elemwise{add,no_inplace} [id G] ''
| | | | |TensorConstant{5} [id D]
| | | | |Subtensor{int64} [id H] ''
| | | | |Shape [id I] ''
| | | | | |Subtensor{:int64:} [id J] ''
| | | | | |<TensorType(int64, vector)> [id K]
| | | | | |ScalarConstant{2} [id L]
| | | | |ScalarConstant{0} [id M]
| | | |Subtensor{:int64:} [id J] ''
| | | |ScalarFromTensor [id N] ''
| | | |Subtensor{int64} [id H] ''
| | |IncSubtensor{Set;:int64:} [id O] ''
| | |AllocEmpty{dtype='int64'} [id P] ''
| | | |Elemwise{add,no_inplace} [id Q] ''
| | | |TensorConstant{5} [id D]
| | | |Subtensor{int64} [id R] ''
| | | |Shape [id S] ''
| | | | |Subtensor{:int64:} [id T] ''
| | | | |<TensorType(int64, vector)> [id U]
| | | | |ScalarConstant{2} [id V]
| | | |ScalarConstant{0} [id W]
| | |Subtensor{:int64:} [id T] ''
| | |ScalarFromTensor [id X] ''
| | |Subtensor{int64} [id R] ''
| |ScalarConstant{2} [id Y]
|Subtensor{int64::} [id Z] ''
|for{cpu,scan_fn}.1 [id C] ''
|ScalarConstant{2} [id BA]
Inner graphs of the scan ops:
for{cpu,scan_fn}.0 [id C] ''
>Elemwise{add,no_inplace} [id BB] ''
> |<TensorType(int64, scalar)> [id BC] -> [id E]
> |<TensorType(int64, scalar)> [id BD] -> [id E]
>Elemwise{add,no_inplace} [id BE] ''
> |<TensorType(int64, scalar)> [id BF] -> [id O]
> |<TensorType(int64, scalar)> [id BG] -> [id O]
for{cpu,scan_fn}.1 [id C] ''
>Elemwise{add,no_inplace} [id BB] ''
>Elemwise{add,no_inplace} [id BE] ''"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
def test_scan_debugprint5():
k = iscalar("k")
A = dvector("A")
# Symbolic description of the result
result, updates = aesara.scan(
fn=lambda prior_result, A: prior_result * A,
outputs_info=aet.ones_like(A),
non_sequences=A,
n_steps=k,
)
final_result = aesara.grad(result[-1].sum(), A)
output_str = debugprint(final_result, file="str")
lines = output_str.split("\n")
expected_output = """Subtensor{int64} [id A] ''
|for{cpu,grad_of_scan_fn}.1 [id B] ''
| |Elemwise{sub,no_inplace} [id C] ''
| | |Subtensor{int64} [id D] ''
| | | |Shape [id E] ''
| | | | |for{cpu,scan_fn} [id F] ''
| | | | |k [id G]
| | | | |IncSubtensor{Set;:int64:} [id H] ''
| | | | | |AllocEmpty{dtype='float64'} [id I] ''
| | | | | | |Elemwise{add,no_inplace} [id J] ''
| | | | | | | |k [id G]
| | | | | | | |Subtensor{int64} [id K] ''
| | | | | | | |Shape [id L] ''
| | | | | | | | |Rebroadcast{0} [id M] ''
| | | | | | | | |InplaceDimShuffle{x,0} [id N] ''
| | | | | | | | |Elemwise{second,no_inplace} [id O] ''
| | | | | | | | |A [id P]
| | | | | | | | |InplaceDimShuffle{x} [id Q] ''
| | | | | | | | |TensorConstant{1.0} [id R]
| | | | | | | |ScalarConstant{0} [id S]
| | | | | | |Subtensor{int64} [id T] ''
| | | | | | |Shape [id U] ''
| | | | | | | |Rebroadcast{0} [id M] ''
| | | | | | |ScalarConstant{1} [id V]
| | | | | |Rebroadcast{0} [id M] ''
| | | | | |ScalarFromTensor [id W] ''
| | | | | |Subtensor{int64} [id K] ''
| | | | |A [id P]
| | | |ScalarConstant{0} [id X]
| | |TensorConstant{1} [id Y]
| |Subtensor{:int64:} [id Z] ''
| | |Subtensor{::int64} [id BA] ''
| | | |Subtensor{:int64:} [id BB] ''
| | | | |for{cpu,scan_fn} [id F] ''
| | | | |ScalarConstant{-1} [id BC]
| | | |ScalarConstant{-1} [id BD]
| | |ScalarFromTensor [id BE] ''
| | |Elemwise{sub,no_inplace} [id C] ''
| |Subtensor{:int64:} [id BF] ''
| | |Subtensor{:int64:} [id BG] ''
| | | |Subtensor{::int64} [id BH] ''
| | | | |for{cpu,scan_fn} [id F] ''
| | | | |ScalarConstant{-1} [id BI]
| | | |ScalarConstant{-1} [id BJ]
| | |ScalarFromTensor [id BK] ''
| | |Elemwise{sub,no_inplace} [id C] ''
| |Subtensor{::int64} [id BL] ''
| | |IncSubtensor{Inc;int64::} [id BM] ''
| | | |Elemwise{second,no_inplace} [id BN] ''
| | | | |for{cpu,scan_fn} [id F] ''
| | | | |InplaceDimShuffle{x,x} [id BO] ''
| | | | |TensorConstant{0.0} [id BP]
| | | |IncSubtensor{Inc;int64} [id BQ] ''
| | | | |Elemwise{second,no_inplace} [id BR] ''
| | | | | |Subtensor{int64::} [id BS] ''
| | | | | | |for{cpu,scan_fn} [id F] ''
| | | | | | |ScalarConstant{1} [id BT]
| | | | | |InplaceDimShuffle{x,x} [id BU] ''
| | | | | |TensorConstant{0.0} [id BV]
| | | | |Elemwise{second} [id BW] ''
| | | | | |Subtensor{int64} [id BX] ''
| | | | | | |Subtensor{int64::} [id BS] ''
| | | | | | |ScalarConstant{-1} [id BY]
| | | | | |InplaceDimShuffle{x} [id BZ] ''
| | | | | |Elemwise{second,no_inplace} [id CA] ''
| | | | | |Sum{acc_dtype=float64} [id CB] ''
| | | | | | |Subtensor{int64} [id BX] ''
| | | | | |TensorConstant{1.0} [id CC]
| | | | |ScalarConstant{-1} [id BY]
| | | |ScalarConstant{1} [id BT]
| | |ScalarConstant{-1} [id CD]
| |Alloc [id CE] ''
| | |TensorConstant{0.0} [id CF]
| | |Elemwise{add,no_inplace} [id CG] ''
| | | |Elemwise{sub,no_inplace} [id C] ''
| | | |TensorConstant{1} [id CH]
| | |Subtensor{int64} [id CI] ''
| | |Shape [id CJ] ''
| | | |A [id P]
| | |ScalarConstant{0} [id CK]
| |A [id P]
|ScalarConstant{-1} [id CL]
Inner graphs of the scan ops:
for{cpu,grad_of_scan_fn}.1 [id B] ''
>Elemwise{add,no_inplace} [id CM] ''
> |Elemwise{mul} [id CN] ''
> | |<TensorType(float64, vector)> [id CO] -> [id BL]
> | |A_copy [id CP] -> [id P]
> |<TensorType(float64, vector)> [id CQ] -> [id BL]
>Elemwise{add,no_inplace} [id CR] ''
> |Elemwise{mul} [id CS] ''
> | |<TensorType(float64, vector)> [id CO] -> [id BL]
> | |<TensorType(float64, vector)> [id CT] -> [id Z]
> |<TensorType(float64, vector)> [id CU] -> [id CE]
for{cpu,scan_fn} [id F] ''
>Elemwise{mul,no_inplace} [id CV] ''
> |<TensorType(float64, vector)> [id CT] -> [id H]
> |A_copy [id CP] -> [id P]
for{cpu,scan_fn} [id F] ''
>Elemwise{mul,no_inplace} [id CV] ''
for{cpu,scan_fn} [id F] ''
>Elemwise{mul,no_inplace} [id CV] ''
for{cpu,scan_fn} [id F] ''
>Elemwise{mul,no_inplace} [id CV] ''
for{cpu,scan_fn} [id F] ''
>Elemwise{mul,no_inplace} [id CV] ''"""
for truth, out in zip(expected_output.split("\n"), lines):
assert truth.strip() == out.strip()
@pytest.mark.skipif(not pydot_imported, reason="pydot not available")
def test_printing_scan():
def f_pow2(x_tm1):
return 2 * x_tm1
state = scalar("state")
n_steps = iscalar("nsteps")
output, updates = aesara.scan(
f_pow2, [], state, [], n_steps=n_steps, truncate_gradient=-1, go_backwards=False
)
f = aesara.function(
[state, n_steps], output, updates=updates, allow_input_downcast=True
)
pydotprint(output, scan_graphs=True)
pydotprint(f, scan_graphs=True)
def test_subtensor():
x = dvector()
y = x[1]
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论