ENH: Allow functions of _period as filter arguments

[buddejul]

Filter functions cannot take arguments that are functions of _period as arguments. Or at least if one other state is provided as argument (see below). Related to #60 which allows _period as an argument.

Examples

def age(_period):
    return _period + 20

# This works
def period_filter(_period):
    return jnp.logical_or(_period > 0, _period < 20)

# This works
def age_filter(age):
    return jnp.logical_or(age > 0, age < 20)

# This doesn't work
def failing_age_filter(age, no_of_kids):
    return jnp.logical_or(age > 0, no_of_kids == 0)

# This works
def workaround_age_filter(_period, no_of_kids):
    return jnp.logical_or(age(_period) > 0, no_of_kids == 0)

When trying to solve with failing_age_filter, lcm returns

ValueError: Not enough or too many arguments provided.