def fn(module, axis):
    if (axis == 0):
        return module.vsplit
    elif (axis == 1):
        return module.hsplit
    else:
        assert (axis == 2)
        return module.dsplit
def fn(module, axis):
    if (axis == 0):
        return module.vsplit
    elif (axis == 1):
        return module.hsplit
    else:
        assert (axis == 2)
        return module.dsplit