(x, y, axis=-1, name=None, op_type="Undefined")
| 1373 | |
| 1374 | |
| 1375 | def _elementwise_op_with_axis(x, y, axis=-1, name=None, op_type="Undefined"): |
| 1376 | assert in_dynamic_or_pir_mode(), ( |
| 1377 | "You can only call `_elementwise_op_with_axis` function within in_dynamic_or_pir_mode" |
| 1378 | ) |
| 1379 | assert op_type in [ |
| 1380 | "add", |
| 1381 | "subtract", |
| 1382 | "multiply", |
| 1383 | "divide", |
| 1384 | ], ( |
| 1385 | f"op_name input error! _elementwise_op_with_axis is an inner function to replace elementwise_add/sub/mul/div. Input op_name={op_type}, Expect op_name=[add|subtract|multiply|divide]\n" |
| 1386 | ) |
| 1387 | op = getattr(_C_ops, op_type) |
| 1388 | x_shape = list(x.shape) |
| 1389 | y_shape = list(y.shape) |
| 1390 | if axis == -1 or len(x_shape) == len(y_shape): |
| 1391 | return op(x, y) |
| 1392 | if len(x_shape) > len(y_shape): |
| 1393 | padding = len(x_shape) - len(y_shape) - axis |
| 1394 | y = paddle.reshape(y, [1] * axis + y_shape + [1] * padding) |
| 1395 | else: |
| 1396 | padding = len(y_shape) - len(x_shape) - axis |
| 1397 | x = paddle.reshape(x, [1] * axis + y_shape + [1] * padding) |
| 1398 | return op(x, y) |
| 1399 | |
| 1400 | |
| 1401 | def _add_with_axis(x, y, axis=-1, name=None): |
no test coverage detected