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]) }, )