def METHOD_NAME(cls, ctx, op):
    (x,), device_id, xp = as_same_device(
        [ctx[c.key] for c in op.inputs], device=op.device, ret_extra=True
    )
    axis1, axis2 = op.axis1, op.axis2
    with device(device_id):
        ctx[op.outputs[0].key] = xp.swapaxes(x, axis1, axis2)