我尝试设计一个数组操作，如下所述：
如果数组中的元素小于0，则新值应设置为0。否则，新值应为原始值的平方。

e.g.
import numpy as np
input = np.array([-2, 5, 2, -1])

# some operation ...

# input should be: [0, 25, 4, 0]

另外，我想让这个操作可以采取任意的输入维度（2d数组，3d数组…）
我只能想到这个解决方案：

input = np.array([-2, 5, 2, -1])
input[input < 0] = 0  # input = [0 5 2 0]

# input[input >= 0] = input*input   <- this will cause error

我已经知道如何完成负片部分，但我不知道如何处理正方形部分，有人能帮我吗？
谢谢！