forward of Split Args: x (CTensor): input tensor. Returns: the output CTensor.
(self, x)
| 4374 | assert self.num_output is not None, "For (parts, num_output), it at least requires one." |
| 4375 | |
| 4376 | def forward(self, x): |
| 4377 | """ |
| 4378 | forward of Split |
| 4379 | Args: |
| 4380 | x (CTensor): input tensor. |
| 4381 | Returns: |
| 4382 | the output CTensor. |
| 4383 | """ |
| 4384 | x_shape = list(x.shape()) |
| 4385 | self.axis = self.axis % len(x_shape) |
| 4386 | if self.parts is None: |
| 4387 | self.parts = [x_shape[self.axis] // self.num_output |
| 4388 | ] * self.num_output |
| 4389 | xs = [] |
| 4390 | _s = 0 |
| 4391 | for _l in self.parts: |
| 4392 | xs.append(singa.SliceOn(x, _s, _s + _l, self.axis)) |
| 4393 | _s += _l |
| 4394 | return tuple(xs) |
| 4395 | |
| 4396 | def backward(self, *dys): |
| 4397 | """ |