[Prim] add meshgrid composite rule #51061
Conversation
|
你的PR提交成功，感谢你对开源项目的贡献! |
|
❌ The PR is not created using PR's template. You can refer to this Demo. |
… add_meshgrid_composite_rule
… add_meshgrid_composite_rule
| """ | ||
| define composite rule of op meshgrid | ||
| If the input has N tensors of size S_0, ... S_n-1, then the output will also have N tensors, where | ||
| each tensor is of shape (S_0, ..., S_n-1) |
There was a problem hiding this comment.
explain diff of each tensor
There was a problem hiding this comment.
done , add an example to demonstrate it
|
|
| if dim == 0: | ||
| shape[i] = 1 | ||
| else: | ||
| shape[i] = inputs[i].shape[0] |
There was a problem hiding this comment.
may the logic be simplified?
There was a problem hiding this comment.
Yes, this op can support zero dim tensor, so we should consider 0 dim tensor as shape [1]
| r2 is Tensor [[4,5], [4,5], [4,5]] | ||
| """ | ||
| size = len(inputs) | ||
| shape = [1 for i in range(size)] |
| shape = [1] * size | ||
| for i in range(size): | ||
| dim = inputs[i].dim() | ||
| assert dim == 0 or dim == 1 |
… add_meshgrid_composite_rule
| @@ -28,6 +28,7 @@ def meshgrid_wrapper(x): | |||
| class TestMeshgridOp(OpTest): | |||
| def setUp(self): | |||
| self.op_type = "meshgrid" | |||
| self.prim_op_type = "comp" | |||
| self.python_api = meshgrid_wrapper | |||
There was a problem hiding this comment.
看一下python3流水线 新增一下public_python_api
PR types
Others
PR changes
Ops
Describe
add composite op meshgrid