Construct new tensors by scanning over axis. Parameters ---------- init: Tensor or list of Tensor The initial condition of first init.shape[0] timestamps update: Tensor or list of Tensor The update rule of the scan given by symbolic tensor. state_placeholder: T
(init, update, state_placeholder, inputs=None, name="scan", tag="", attrs=None)
| 140 | |
| 141 | |
| 142 | def scan(init, update, state_placeholder, inputs=None, name="scan", tag="", attrs=None): |
| 143 | """Construct new tensors by scanning over axis. |
| 144 | |
| 145 | Parameters |
| 146 | ---------- |
| 147 | init: Tensor or list of Tensor |
| 148 | The initial condition of first init.shape[0] timestamps |
| 149 | |
| 150 | update: Tensor or list of Tensor |
| 151 | The update rule of the scan given by symbolic tensor. |
| 152 | |
| 153 | state_placeholder: Tensor or list of Tensor |
| 154 | The placeholder variables used by update. |
| 155 | |
| 156 | inputs: Tensor or list of Tensor, optional |
| 157 | The list of inputs to the scan. This is not required, but can |
| 158 | be useful for the compiler to detect scan body faster. |
| 159 | |
| 160 | name: str, optional |
| 161 | The name hint of the tensor |
| 162 | |
| 163 | tag: str, optional |
| 164 | Additonal tag information about the compute. |
| 165 | |
| 166 | attrs: dict, optional |
| 167 | The additional auxiliary attributes about the compute. |
| 168 | |
| 169 | Returns |
| 170 | ------- |
| 171 | tensor: Tensor or list of Tensors |
| 172 | The created tensor or tuple of tensors contains multiple outputs. |
| 173 | |
| 174 | Example |
| 175 | ------- |
| 176 | .. code-block:: python |
| 177 | |
| 178 | # The following code is equivalent to numpy.cumsum |
| 179 | m = te.var("m") |
| 180 | n = te.var("n") |
| 181 | X = te.placeholder((m, n), name="X") |
| 182 | s_state = te.placeholder((m, n)) |
| 183 | s_init = te.compute((1, n), lambda _, i: X[0, i]) |
| 184 | s_update = te.compute((m, n), lambda t, i: s_state[t-1, i] + X[t, i]) |
| 185 | res = tvm.te.scan(s_init, s_update, s_state, X) |
| 186 | """ |
| 187 | if _tag.TagScope.get_current() is not None: |
| 188 | if tag != "": |
| 189 | raise ValueError("nested tag is not allowed for now") |
| 190 | tag = _tag.TagScope.get_current().tag |
| 191 | if isinstance(init, _tensor.Tensor): |
| 192 | init = [init] |
| 193 | if isinstance(update, _tensor.Tensor): |
| 194 | update = [update] |
| 195 | if isinstance(state_placeholder, _tensor.Tensor): |
| 196 | state_placeholder = [state_placeholder] |
| 197 | if isinstance(inputs, _tensor.Tensor): |
| 198 | inputs = [inputs] |
| 199 | if inputs is None: |
nothing calls this directly
no test coverage detected
searching dependent graphs…