Skip to content

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

Split ScanInfo.tap_array into mit-mot, mit-sot, sit-sot parts

上级 cdd8575b
隐藏空白字符变更
内嵌 并排
正在显示 包含 174 行增加 和 138 行删除
aesara/scan/basic.py
浏览文件 @ 76585fe1
......@@ -1123,15 +1123,14 @@ def scan(
# Step 7. Create the Scan Op
##
tap_array = tuple(tuple(v) for v in mit_sot_tap_array) + tuple(
(-1,) for x in range(n_sit_sot)
)
if allow_gc is None:
allow_gc = config.scan__allow_gc
info = ScanInfo(
tap_array=tap_array,
n_seqs=n_seqs,
mit_mot_in_slices=(),
mit_sot_in_slices=tuple(tuple(v) for v in mit_sot_tap_array),
sit_sot_in_slices=tuple((-1,) for x in range(n_sit_sot)),
n_mit_mot=n_mit_mot,
n_mit_mot_outs=n_mit_mot_outs,
mit_mot_out_slices=tuple(tuple(v) for v in mit_mot_out_slices),
......
aesara/scan/op.py
浏览文件 @ 76585fe1
......@@ -47,7 +47,7 @@ import dataclasses
import logging
import time
from collections import OrderedDict
from itertools import product
from itertools import chain, product
from typing import Callable, List, Optional, Union
import numpy as np
......@@ -203,11 +203,9 @@ def copy_var_format(var, as_var):
@dataclasses.dataclass(frozen=True)
class ScanInfo:
tap_array: tuple
"""
This is a tuple containing tuples of inner-output lag/lead values for the
mit-mots, mit-sots, and ``[-1]`` for each sit-sot.
"""
mit_mot_in_slices: tuple
mit_sot_in_slices: tuple
sit_sot_in_slices: tuple
n_seqs: int
n_mit_mot: int
n_mit_mot_outs: int
......@@ -219,6 +217,10 @@ class ScanInfo:
n_non_seqs: int
as_while: bool
@property
def tap_array(self):
return self.mit_mot_in_slices + self.mit_sot_in_slices + self.sit_sot_in_slices
TensorConstructorType = Callable[[List[bool], Union[str, np.generic]], TensorType]
......@@ -235,7 +237,7 @@ class ScanMethodsMixin:
return list_inputs[1 : 1 + self.info.n_seqs]
def inner_mitmot(self, list_inputs):
n_taps = sum(len(x) for x in self.info.tap_array[: self.info.n_mit_mot])
n_taps = sum(len(x) for x in self.info.mit_mot_in_slices)
return list_inputs[self.info.n_seqs : self.info.n_seqs + n_taps]
def outer_mitmot(self, list_inputs):
......@@ -251,16 +253,15 @@ class ScanMethodsMixin:
return list_outputs[: self.info.n_mit_mot]
def mitmot_taps(self):
return self.info.tap_array[: self.info.n_mit_mot]
return self.info.mit_mot_in_slices
def mitmot_out_taps(self):
return self.info.mit_mot_out_slices[: self.info.n_mit_mot]
def inner_mitsot(self, list_inputs):
n_mitmot_taps = sum(len(x) for x in self.info.tap_array[: self.info.n_mit_mot])
ntaps_upto_sit_sot = sum(
len(x)
for x in self.info.tap_array[: (self.info.n_mit_mot + self.info.n_mit_sot)]
n_mitmot_taps = sum(len(x) for x in self.info.mit_mot_in_slices)
ntaps_upto_sit_sot = n_mitmot_taps + sum(
len(x) for x in self.info.mit_sot_in_slices
)
return list_inputs[
self.info.n_seqs + n_mitmot_taps : self.info.n_seqs + ntaps_upto_sit_sot
......@@ -280,14 +281,12 @@ class ScanMethodsMixin:
]
def mitsot_taps(self):
return self.info.tap_array[
self.info.n_mit_mot : self.info.n_mit_mot + self.info.n_mit_sot
]
return self.info.mit_sot_in_slices
def inner_sitsot(self, list_inputs):
n_taps_upto_sit_sot = sum(
len(x)
for x in self.info.tap_array[: (self.info.n_mit_mot + self.info.n_mit_sot)]
for x in chain(self.info.mit_mot_in_slices, self.info.mit_sot_in_slices)
)
offset = self.info.n_seqs + n_taps_upto_sit_sot
return list_inputs[offset : offset + self.info.n_sit_sot]
......@@ -328,7 +327,7 @@ class ScanMethodsMixin:
def inner_shared(self, list_inputs):
n_taps_upto_sit_sot = sum(
len(x)
for x in self.info.tap_array[: (self.info.n_mit_mot + self.info.n_mit_sot)]
for x in chain(self.info.mit_mot_in_slices, self.info.mit_sot_in_slices)
)
offset = self.info.n_seqs + n_taps_upto_sit_sot + self.info.n_sit_sot
return list_inputs[offset : offset + self.info.n_shared_outs]
......@@ -362,7 +361,7 @@ class ScanMethodsMixin:
def inner_non_seqs(self, list_inputs):
n_taps_upto_sit_sot = sum(
len(x)
for x in self.info.tap_array[: (self.info.n_mit_mot + self.info.n_mit_sot)]
for x in chain(self.info.mit_mot_in_slices, self.info.mit_sot_in_slices)
)
offset = (
self.info.n_seqs
......@@ -427,8 +426,14 @@ class ScanMethodsMixin:
outer_oidx += 0
# Handle mitmots, mitsots and sitsots variables
for i in range(len(self.info.tap_array)):
nb_input_taps = len(self.info.tap_array[i])
for i, tap in enumerate(
chain(
self.info.mit_mot_in_slices,
self.info.mit_sot_in_slices,
self.info.sit_sot_in_slices,
)
):
nb_input_taps = len(tap)
if i < self.info.n_mit_mot:
nb_output_taps = len(self.info.mit_mot_out_slices[i])
......@@ -758,7 +763,12 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
self.name = "scan_fn"
# Pre-computing some values to speed up perform
self.mintaps = [np.min(x) for x in info.tap_array]
self.mintaps = [
min(x)
for x in chain(
info.mit_mot_in_slices, info.mit_sot_in_slices, info.sit_sot_in_slices
)
]
self.mintaps += [0 for x in range(info.n_nit_sot)]
self.seqs_arg_offset = 1 + info.n_seqs
self.shared_arg_offset = (
......@@ -813,7 +823,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
info = self.info
input_idx = info.n_seqs
for mitmot_idx in range(info.n_mit_mot):
for inp_tap in info.tap_array[mitmot_idx]:
for inp_tap in info.mit_mot_in_slices[mitmot_idx]:
if inp_tap in info.mit_mot_out_slices[mitmot_idx]:
# Figure out the index of the corresponding output
output_idx = sum(
......@@ -1292,7 +1302,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
input_idx = self.info.n_seqs
for mitmot_idx in range(self.info.n_mit_mot):
for inp_tap in self.info.tap_array[mitmot_idx]:
for inp_tap in self.info.mit_mot_in_slices[mitmot_idx]:
if inp_tap in self.info.mit_mot_out_slices[mitmot_idx]:
inp = self.inputs[input_idx]
......@@ -1447,7 +1457,14 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
cython_mintaps = np.asarray(self.mintaps, dtype="int32")
tap_array_len = tuple(len(x) for x in self.info.tap_array)
tap_array_len = tuple(
len(x)
for x in chain(
self.info.mit_mot_in_slices,
self.info.mit_sot_in_slices,
self.info.sit_sot_in_slices,
)
)
cython_vector_seqs = np.asarray(self.vector_seqs, dtype="int32")
cython_vector_outs = np.asarray(self.vector_outs, dtype="int32")
......@@ -1506,7 +1523,9 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
self.info.n_nit_sot,
self.info.as_while,
cython_mintaps,
self.info.tap_array,
self.info.mit_mot_in_slices
+ self.info.mit_sot_in_slices
+ self.info.sit_sot_in_slices,
tap_array_len,
cython_vector_seqs,
cython_vector_outs,
......@@ -1679,7 +1698,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
offset = self.nit_sot_arg_offset + info.n_nit_sot
other_args = inputs[offset:]
inner_input_storage = self.fn.input_storage
nb_mitmot_in = sum(map(len, info.tap_array[: info.n_mit_mot]))
nb_mitmot_in = sum(map(len, info.mit_mot_in_slices))
old_mitmot_input_storage = [None] * nb_mitmot_in
old_mitmot_input_data = [None] * nb_mitmot_in
inner_output_storage = self.fn.output_storage
......@@ -1688,7 +1707,14 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
fn = self.fn.fn
offset = (
info.n_seqs
+ sum(map(len, info.tap_array[: self.n_outs]))
+ sum(
len(x)
for x in chain(
info.mit_mot_in_slices,
info.mit_sot_in_slices,
info.sit_sot_in_slices,
)
)
+ info.n_shared_outs
)
for idx in range(len(other_args)):
......@@ -1710,17 +1736,23 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
inner_input_storage[idx].storage[0] = seqs[idx][i]
offset = info.n_seqs
for idx in range(self.n_outs):
for idx, taps in enumerate(
chain(
info.mit_mot_in_slices,
info.mit_sot_in_slices,
info.sit_sot_in_slices,
)
):
if self.vector_outs[idx]:
for tap in info.tap_array[idx]:
_idx = (pos[idx] + tap) % store_steps[idx]
for t in taps:
_idx = (pos[idx] + t) % store_steps[idx]
inner_input_storage[offset].storage[0] = output_storage[idx][0][
_idx : _idx + 1
].reshape(())
offset += 1
else:
for tap in info.tap_array[idx]:
_idx = (pos[idx] + tap) % store_steps[idx]
for t in taps:
_idx = (pos[idx] + t) % store_steps[idx]
inner_input_storage[offset].storage[0] = output_storage[idx][0][
_idx
]
......@@ -1864,11 +1896,11 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
# 5.3 Copy over the values for mit_mot outputs
mitmot_inp_offset = 0
mitmot_out_idx = 0
for j in range(info.n_mit_mot):
for j, taps in enumerate(info.mit_mot_in_slices):
for k in info.mit_mot_out_slices[j]:
if self.mitmots_preallocated[mitmot_out_idx]:
# This output tap has been preallocated.
inp_idx = mitmot_inp_offset + info.tap_array[j].index(k)
inp_idx = mitmot_inp_offset + taps.index(k)
# Verify whether the input points to the same data as
# it did before the execution of the inner function.
......@@ -1899,7 +1931,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
mitmot_out_idx += 1
mitmot_inp_offset += len(info.tap_array[j])
mitmot_inp_offset += len(taps)
# 5.4 Copy over the values for mit_sot/sit_sot outputs
begin = info.n_mit_mot
......@@ -2141,14 +2173,17 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
n_outs = info.n_mit_mot + info.n_mit_sot + info.n_sit_sot
outs_shape = []
for idx in range(n_outs):
abs(min(info.tap_array[idx]))
for k in info.tap_array[idx]:
for idx, taps in enumerate(
chain(
info.mit_mot_in_slices, info.mit_sot_in_slices, info.sit_sot_in_slices
)
):
for k in taps:
outs_shape += [input_shapes[idx + info.n_seqs + 1][1:]]
# if extra_infer_shape:
# mintap = abs(min(taps))
# corresponding_tap = node.inputs[outer_inp_idx][mintap + k]
# out_equivalent[self.inputs[inner_inp_idx]] = corresponding_tap
# inner_inp_idx += 1
outer_inp_idx += 1
# shared_outs
......@@ -2511,7 +2546,14 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
# Check if the pos-th input is associated with one of the
# recurrent states
x_is_state = pos < sum(len(t) for t in info.tap_array)
x_is_state = pos < sum(
len(t)
for t in chain(
info.mit_mot_in_slices,
info.mit_sot_in_slices,
info.sit_sot_in_slices,
)
)
if x_is_state and len(idxs) > 0:
opos = idxs[0]
......@@ -2537,14 +2579,16 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
outer_inp_seqs = [x[n_steps - 1 :: -1] for x in inputs[1 : 1 + info.n_seqs]]
else:
outer_inp_seqs = [x[::-1] for x in inputs[1 : 1 + info.n_seqs]]
for idx in range(info.n_mit_mot + info.n_mit_sot):
mintap = np.min(info.tap_array[idx])
for idx, taps in enumerate(
chain(info.mit_mot_in_slices, info.mit_sot_in_slices)
):
mintap = min(taps)
if idx < info.n_mit_mot:
outmaxtap = np.max(self.mitmot_out_taps()[idx])
else:
outmaxtap = 0
seq = outs[idx]
for k in info.tap_array[idx]:
for k in taps:
if outmaxtap - k != 0:
nw_seq = seq[k - mintap : -(outmaxtap - k)][::-1]
else:
......@@ -2611,7 +2655,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
n_mitmot_outs = 0
n_mitmot_inps = 0
for idx in range(info.n_mit_mot):
for idx, taps in enumerate(info.mit_mot_in_slices):
if isinstance(dC_douts[idx].type, DisconnectedType):
out = outs[idx]
outer_inp_mitmot.append(at.zeros_like(out))
......@@ -2623,14 +2667,14 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
through_shared = False
disconnected = True
for jdx in range(len(info.mit_mot_out_slices[idx])):
for mit_mot_out_slice in info.mit_mot_out_slices[idx]:
inner_inp_mitmot.append(dC_dXts[out_pos])
mitmot_inp_taps[idx].append(-info.mit_mot_out_slices[idx][jdx])
mitmot_inp_taps[idx].append(-mit_mot_out_slice)
n_mitmot_inps += 1
out_pos += 1
for jdx in range(len(info.tap_array[idx])):
tap = -info.tap_array[idx][jdx]
for tap in taps:
tap = -tap
# Only create a new inner input if there is not already one
# associated with this input tap
......@@ -2656,29 +2700,27 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
idx
].index(tap)
replacement = inner_inp_mitmot[-replacement_idx]
info.tap_array[idx]
new_inner_out_mitmot = clone_replace(
new_inner_out_mitmot, replace=[(to_replace, replacement)]
)
inner_out_mitmot.append(new_inner_out_mitmot)
if not disconnected_dC_dinps_t[ins_pos]:
disconnected = False
disconnected &= disconnected_dC_dinps_t[ins_pos]
for _sh in self.inner_shared(self_inputs):
if _sh in graph_inputs([dC_dinps_t[ins_pos]]):
through_shared = True
through_shared = any(
_sh in graph_inputs([dC_dinps_t[ins_pos]])
for _sh in self.inner_shared(self_inputs)
)
ins_pos += 1
n_mitmot_outs += 1
mitmot_out_taps[idx].append(-info.tap_array[idx][jdx])
mitmot_out_taps[idx].append(tap)
# Only add the tap as a new input tap if needed
if tap not in mitmot_inp_taps[idx]:
n_mitmot_inps += 1
mitmot_inp_taps[idx].append(-info.tap_array[idx][jdx])
mitmot_inp_taps[idx].append(tap)
if undefined_msg:
type_outs.append(undefined_msg)
......@@ -2690,14 +2732,13 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
type_outs.append("connected")
offset = info.n_mit_mot
for idx in range(info.n_mit_sot):
for idx, taps in enumerate(info.mit_sot_in_slices):
if isinstance(dC_douts[idx + offset].type, DisconnectedType):
outer_inp_mitmot.append(outs[idx + offset].zeros_like())
else:
outer_inp_mitmot.append(dC_douts[idx + offset][::-1])
mitmot_inp_taps.append([])
mitmot_out_taps.append([])
idx_tap = idx + info.n_mit_mot
inner_inp_mitmot.append(dC_dXts[out_pos])
out_pos += 1
n_mitmot_inps += 1
......@@ -2705,7 +2746,8 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
through_shared = False
disconnected = True
mitmot_inp_taps[idx + offset].append(0)
for jdx in range(len(info.tap_array[idx_tap])):
for tap in taps:
tap = -tap
inner_inp_mitmot.append(dC_dXtm1s[ins_pos - info.n_seqs])
if isinstance(dC_dinps_t[ins_pos].type, NullType):
......@@ -2718,13 +2760,15 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
else:
inner_out_mitmot.append(dC_dinps_t[ins_pos])
mitmot_inp_taps[idx + offset].append(-info.tap_array[idx_tap][jdx])
mitmot_out_taps[idx].append(-info.tap_array[idx_tap][jdx])
if not disconnected_dC_dinps_t[ins_pos]:
disconnected = False
for _sh in self.inner_shared(self_inputs):
if _sh in graph_inputs([dC_dinps_t[ins_pos]]):
through_shared = True
mitmot_inp_taps[idx + offset].append(tap)
mitmot_out_taps[idx].append(tap)
disconnected &= disconnected_dC_dinps_t[ins_pos]
through_shared = any(
_sh in graph_inputs([dC_dinps_t[ins_pos]])
for _sh in self.inner_shared(self_inputs)
)
n_mitmot_inps += 1
ins_pos += 1
......@@ -2770,9 +2814,10 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
else:
inner_out_mitmot.append(dC_dinps_t[ins_pos])
for _sh in self.inner_shared(self_inputs):
if _sh in graph_inputs([dC_dinps_t[ins_pos]]):
through_shared = True
through_shared = any(
_sh in graph_inputs([dC_dinps_t[ins_pos]])
for _sh in self.inner_shared(self_inputs)
)
if isinstance(dC_dinps_t[ins_pos].type, NullType):
type_outs.append(dC_dinps_t[ins_pos].type.why_null)
......@@ -2860,7 +2905,6 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
)
n_sitsot_outs = len(outer_inp_sitsot)
new_tap_array = mitmot_inp_taps + [[-1] for k in range(n_sitsot_outs)]
outer_inputs = (
[grad_steps]
......@@ -2884,7 +2928,9 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
out_info = ScanInfo(
n_seqs=len(outer_inp_seqs),
n_mit_sot=0,
tap_array=tuple(tuple(v) for v in new_tap_array),
mit_mot_in_slices=tuple(tuple(v) for v in mitmot_inp_taps),
mit_sot_in_slices=(),
sit_sot_in_slices=tuple((-1,) for k in range(n_sitsot_outs)),
n_mit_mot=len(outer_inp_mitmot),
n_mit_mot_outs=n_mitmot_outs,
mit_mot_out_slices=tuple(tuple(v) for v in mitmot_out_taps),
......@@ -3065,16 +3111,9 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
# evan point for the number of nit_sot which I think should just be
# ignored (?)
new_tap_array = []
b = 0
e = info.n_mit_mot
new_tap_array += info.tap_array[b:e] * 2
b = e
e += info.n_mit_sot
new_tap_array += info.tap_array[b:e] * 2
b = e
e += info.n_sit_sot
new_tap_array += info.tap_array[b:e] * 2
new_mit_mot_in_slices = info.mit_mot_in_slices * 2
new_mit_sot_in_slices = info.mit_sot_in_slices * 2
new_sit_sot_in_slices = info.sit_sot_in_slices * 2
# Sequences ...
b = 1
......@@ -3095,7 +3134,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
b = e
e = e + info.n_mit_mot
ib = ie
ie = ie + int(sum(len(x) for x in info.tap_array[: info.n_mit_mot]))
ie = ie + int(sum(len(x) for x in info.mit_mot_in_slices))
clean_eval_points = []
for inp, evp in zip(inputs[b:e], eval_points[b:e]):
if evp is not None:
......@@ -3110,14 +3149,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
b = e
e = e + info.n_mit_sot
ib = ie
ie = ie + int(
sum(
len(x)
for x in info.tap_array[
info.n_mit_mot : info.n_mit_mot + info.n_mit_sot
]
)
)
ie = ie + int(sum(len(x) for x in info.mit_sot_in_slices))
clean_eval_points = []
for inp, evp in zip(inputs[b:e], eval_points[b:e]):
if evp is not None:
......@@ -3217,13 +3249,15 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
out_info = ScanInfo(
n_seqs=info.n_seqs * 2,
mit_mot_in_slices=new_mit_mot_in_slices,
mit_sot_in_slices=new_mit_sot_in_slices,
sit_sot_in_slices=new_sit_sot_in_slices,
n_mit_sot=info.n_mit_sot * 2,
n_sit_sot=info.n_sit_sot * 2,
n_mit_mot=info.n_mit_mot * 2,
n_nit_sot=info.n_nit_sot * 2,
n_shared_outs=info.n_shared_outs,
n_mit_mot_outs=n_mit_mot_outs * 2,
tap_array=tuple(tuple(v) for v in new_tap_array),
mit_mot_out_slices=tuple(tuple(v) for v in info.mit_mot_out_slices) * 2,
n_non_seqs=len(inner_other),
as_while=info.as_while,
......
aesara/scan/opt.py
浏览文件 @ 76585fe1
......@@ -2,6 +2,7 @@
import copy
import dataclasses
from itertools import chain
from sys import maxsize
from typing import Dict, List, Optional, Tuple
......@@ -82,7 +83,7 @@ def remove_constants_and_unused_inputs_scan(fgraph, node):
op = node.op
# We only need to take care of sequences and other arguments
st = op.n_seqs
st += int(sum(len(x) for x in op.tap_array[: (op.n_mit_mot + op.n_mit_sot)]))
st += int(sum(len(x) for x in chain(op.mit_mot_in_slices, op.mit_sot_in_slices)))
st += op.n_sit_sot
st += op.n_shared_outs
......@@ -1147,7 +1148,7 @@ def save_mem_new_scan(fgraph, node):
c_outs = op.n_mit_mot + op.n_mit_sot + op.n_sit_sot + op.n_nit_sot
init_l = [0 for x in range(op.n_mit_mot)]
init_l += [abs(min(v)) for v in op.tap_array[op.n_mit_mot :]]
init_l += [abs(min(v)) for v in chain(op.mit_sot_in_slices, op.sit_sot_in_slices)]
init_l += [0 for x in range(op.n_nit_sot)]
# 2. Check the clients of each output and see for how many steps
# does scan need to run
......@@ -1678,30 +1679,29 @@ class ScanMerge(GlobalOptimizer):
inner_ins[idx].append(rename(nd.op.inner_seqs(nd.op.inputs), idx))
outer_ins += rename(nd.op.outer_seqs(nd.inputs), idx)
tap_array = ()
mit_mot_out_slices = ()
mit_mot_in_slices = ()
for idx, nd in enumerate(nodes):
# MitMot
inner_ins[idx].append(rename(nd.op.inner_mitmot(nd.op.inputs), idx))
inner_outs[idx].append(nd.op.inner_mitmot_outs(nd.op.outputs))
tap_array += nd.op.mitmot_taps()
mit_mot_in_slices += nd.op.mitmot_taps()
mit_mot_out_slices += nd.op.mitmot_out_taps()
outer_ins += rename(nd.op.outer_mitmot(nd.inputs), idx)
outer_outs += nd.op.outer_mitmot_outs(nd.outputs)
mit_sot_in_slices = ()
for idx, nd in enumerate(nodes):
# MitSot
inner_ins[idx].append(rename(nd.op.inner_mitsot(nd.op.inputs), idx))
inner_outs[idx].append(nd.op.inner_mitsot_outs(nd.op.outputs))
tap_array += nd.op.mitsot_taps()
mit_sot_in_slices += nd.op.mitsot_taps()
outer_ins += rename(nd.op.outer_mitsot(nd.inputs), idx)
outer_outs += nd.op.outer_mitsot_outs(nd.outputs)
sit_sot_in_slices = ()
for idx, nd in enumerate(nodes):
# SitSot
inner_ins[idx].append(rename(nd.op.inner_sitsot(nd.op.inputs), idx))
tap_array += tuple((-1,) for x in range(nd.op.n_sit_sot))
sit_sot_in_slices += tuple((-1,) for x in range(nd.op.n_sit_sot))
inner_outs[idx].append(nd.op.inner_sitsot_outs(nd.op.outputs))
outer_ins += rename(nd.op.outer_sitsot(nd.inputs), idx)
outer_outs += nd.op.outer_sitsot_outs(nd.outputs)
......@@ -1794,7 +1794,9 @@ class ScanMerge(GlobalOptimizer):
new_inner_outs += inner_outs[idx][gr_idx]
info = ScanInfo(
tap_array=tap_array,
mit_mot_in_slices=mit_mot_in_slices,
mit_sot_in_slices=mit_sot_in_slices,
sit_sot_in_slices=sit_sot_in_slices,
n_seqs=sum(nd.op.n_seqs for nd in nodes),
n_mit_mot=sum(nd.op.n_mit_mot for nd in nodes),
n_mit_mot_outs=sum(nd.op.n_mit_mot_outs for nd in nodes),
......@@ -2212,14 +2214,10 @@ def push_out_dot1_scan(fgraph, node):
inner_non_seqs = op.inner_non_seqs(op.inputs)
outer_non_seqs = op.outer_non_seqs(node.inputs)
st = len(op.mitmot_taps()) + len(op.mitsot_taps())
new_info = dataclasses.replace(
op.info,
tap_array=(
op.info.tap_array[: st + idx]
+ op.info.tap_array[st + idx + 1 :]
),
sit_sot_in_slices=op.info.sit_sot_in_slices[:idx]
+ op.info.sit_sot_in_slices[idx + 1 :],
n_sit_sot=op.info.n_sit_sot - 1,
n_nit_sot=op.info.n_nit_sot + 1,
)
......
aesara/scan/utils.py
浏览文件 @ 76585fe1
......@@ -4,6 +4,7 @@ import copy
import dataclasses
import logging
from collections import OrderedDict, namedtuple
from itertools import chain
from typing import TYPE_CHECKING, Callable, List, Optional, Sequence, Set, Tuple, Union
from typing import cast as type_cast
......@@ -348,11 +349,12 @@ class Validator:
def scan_can_remove_outs(op, out_idxs):
"""
Looks at all outputs defined by indices ``out_idxs`` and see whom can be
removed from the scan op without affecting the rest. Return two lists,
the first one with the indices of outs that can be removed, the second
with the outputs that can not be removed.
"""Look at all outputs defined by indices ``out_idxs`` and determines which can be removed.
Returns
-------
two lists, the first one with the indices of outs that can be removed, the
second with the outputs that can not be removed.
"""
non_removable = [o for i, o in enumerate(op.outputs) if i not in out_idxs]
......@@ -360,9 +362,10 @@ def scan_can_remove_outs(op, out_idxs):
out_ins = []
offset = op.n_seqs
lim = op.n_mit_mot + op.n_mit_sot + op.n_sit_sot
for idx in range(lim):
n_ins = len(op.info.tap_array[idx])
for idx, tap in enumerate(
chain(op.mit_mot_in_slices, op.mit_sot_in_slices, op.sit_sot_in_slices)
):
n_ins = len(tap)
out_ins += [op.inputs[offset : offset + n_ins]]
offset += n_ins
out_ins += [[] for k in range(op.n_nit_sot)]
......@@ -398,7 +401,9 @@ def compress_outs(op, not_required, inputs):
from aesara.scan.op import ScanInfo
info = ScanInfo(
tap_array=(),
mit_mot_in_slices=(),
mit_sot_in_slices=(),
sit_sot_in_slices=(),
n_seqs=op.info.n_seqs,
n_mit_mot=0,
n_mit_mot_outs=0,
......@@ -428,23 +433,24 @@ def compress_outs(op, not_required, inputs):
info = dataclasses.replace(
info,
n_mit_mot=info.n_mit_mot + 1,
tap_array=info.tap_array + (tuple(op.tap_array[offset + idx]),),
mit_mot_in_slices=info.mit_mot_in_slices + (op.mit_mot_in_slices[idx],),
mit_mot_out_slices=info.mit_mot_out_slices
+ (tuple(op.mit_mot_out_slices[offset + idx]),),
+ (op.mit_mot_out_slices[idx],),
)
# input taps
for jdx in op.tap_array[offset + idx]:
for jdx in op.mit_mot_in_slices[idx]:
op_inputs += [op.inputs[i_offset]]
i_offset += 1
# output taps
for jdx in op.mit_mot_out_slices[offset + idx]:
for jdx in op.mit_mot_out_slices[idx]:
op_outputs += [op.outputs[o_offset]]
o_offset += 1
# node inputs
node_inputs += [inputs[ni_offset + idx]]
else:
o_offset += len(op.mit_mot_out_slices[offset + idx])
i_offset += len(op.tap_array[offset + idx])
o_offset += len(op.mit_mot_out_slices[idx])
i_offset += len(op.mit_mot_in_slices[idx])
info = dataclasses.replace(info, n_mit_mot_outs=len(op_outputs))
offset += op.n_mit_mot
ni_offset += op.n_mit_mot
......@@ -456,10 +462,10 @@ def compress_outs(op, not_required, inputs):
info = dataclasses.replace(
info,
n_mit_sot=info.n_mit_sot + 1,
tap_array=info.tap_array + (tuple(op.tap_array[offset + idx]),),
mit_sot_in_slices=info.mit_sot_in_slices + (op.mit_sot_in_slices[idx],),
)
# input taps
for jdx in op.tap_array[offset + idx]:
for jdx in op.mit_sot_in_slices[idx]:
op_inputs += [op.inputs[i_offset]]
i_offset += 1
# output taps
......@@ -469,7 +475,7 @@ def compress_outs(op, not_required, inputs):
node_inputs += [inputs[ni_offset + idx]]
else:
o_offset += 1
i_offset += len(op.tap_array[offset + idx])
i_offset += len(op.mit_sot_in_slices[idx])
offset += op.n_mit_sot
ni_offset += op.n_mit_sot
......@@ -480,7 +486,7 @@ def compress_outs(op, not_required, inputs):
info = dataclasses.replace(
info,
n_sit_sot=info.n_sit_sot + 1,
tap_array=info.tap_array + (tuple(op.tap_array[offset + idx]),),
sit_sot_in_slices=info.sit_sot_in_slices + (op.sit_sot_in_slices[idx],),
)
# input taps
op_inputs += [op.inputs[i_offset]]
......@@ -613,8 +619,9 @@ class ScanArgs:
n_mit_mot = info.n_mit_mot
n_mit_sot = info.n_mit_sot
self.mit_mot_in_slices = list(info.tap_array[:n_mit_mot])
self.mit_sot_in_slices = list(info.tap_array[n_mit_mot : n_mit_mot + n_mit_sot])
self.mit_mot_in_slices = info.mit_mot_in_slices
self.mit_sot_in_slices = info.mit_sot_in_slices
self.sit_sot_in_slices = info.sit_sot_in_slices
n_mit_mot_ins = sum(len(s) for s in self.mit_mot_in_slices)
n_mit_sot_ins = sum(len(s) for s in self.mit_sot_in_slices)
......@@ -800,14 +807,12 @@ class ScanArgs:
from aesara.scan.op import ScanInfo
return ScanInfo(
mit_mot_in_slices=tuple(tuple(v) for v in self.mit_mot_in_slices),
mit_sot_in_slices=tuple(tuple(v) for v in self.mit_sot_in_slices),
sit_sot_in_slices=((-1,),) * len(self.inner_in_sit_sot),
n_seqs=len(self.outer_in_seqs),
n_mit_mot=len(self.outer_in_mit_mot),
n_mit_sot=len(self.outer_in_mit_sot),
tap_array=(
tuple(tuple(v) for v in self.mit_mot_in_slices)
+ tuple(tuple(v) for v in self.mit_sot_in_slices)
+ ((-1,),) * len(self.inner_in_sit_sot)
),
n_sit_sot=len(self.outer_in_sit_sot),
n_nit_sot=len(self.outer_in_nit_sot),
n_shared_outs=len(self.outer_in_shared),
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论