Takes a tensor as input and outputs a tensor containing the shape of the input tensor.
| 3353 | |
| 3354 | |
| 3355 | class Shape(Operator): |
| 3356 | """ |
| 3357 | Takes a tensor as input and outputs a tensor containing the shape of the |
| 3358 | input tensor. |
| 3359 | """ |
| 3360 | |
| 3361 | def __init__(self): |
| 3362 | super(Shape, self).__init__() |
| 3363 | |
| 3364 | def forward(self, x): |
| 3365 | """ |
| 3366 | Args: |
| 3367 | x (CTensor): Input tensor |
| 3368 | Returns: |
| 3369 | CTensor, the output |
| 3370 | """ |
| 3371 | cur = list(x.shape()) |
| 3372 | cur = tensor.from_numpy(np.array(cur)) |
| 3373 | cur.to_device(x.device()) |
| 3374 | return cur.data |
| 3375 | |
| 3376 | def backward(self, dy): |
| 3377 | """ |
| 3378 | Args: |
| 3379 | dy (CTensor): the gradient tensor from upper operations |
| 3380 | Returns: |
| 3381 | list of int, the shape of dy |
| 3382 | """ |
| 3383 | return list(dy.shape()) |
| 3384 | |
| 3385 | |
| 3386 | def shape(x): |
no outgoing calls
no test coverage detected