def template(x_order=OrderNC,
y_order=OrderNC,
axis=Axis.C,
description: str = ""):
shape = (np.arange(x_order.ndim) + 2).tolist()
vx = np.random.rand(*shape) - 0.5
vy = np.exp(vx) / np.sum(
np.exp(vx), axis=x_order.axes_dict[axis], keepdims=True)
x = Variable(vx.shape, order=x_order)
y, = Softmax(None, axis=axis)(x)
y.change_order(y_order)
generate_kernel_test_case(
description=f"Softmax {description}",
graph=Graph([x], [y]),
inputs={x: vx},
backend=["webgpu", "webassembly"],
expected={
y: np.transpose(vy, [x_order.axes_dict[a] for a in y.order.axes])
},
)
def template(x_order=OrderNC,
y_order=OrderNC,
axis=Axis.C,
description: str = ""):
vx = np.random.rand(2, 3) - 0.5
vy = np.exp(vx) / np.sum(
np.exp(vx), axis=OrderNC.axes_dict[axis], keepdims=True)
x = Variable(vx.shape, order=OrderNC)
y, = Softmax(None, axis=axis)(x)
x.change_order(x_order)
y.change_order(y_order)
generate_kernel_test_case(
description=f"Softmax {description}",
graph=Graph([x], [y]),
inputs={
x: np.transpose(vx, [OrderNC.axes_dict[a] for a in x.order.axes])
},
expected={
y: np.transpose(vy, [OrderNC.axes_dict[a] for a in y.order.axes])
},
)
