Update functools.py

[timmens]

In this PR, I

• Add type hints to functools.py
• Replace the allow_kwargs decorator with the allow_only_kwargs decorator
• Open issue for further improvements of decorator return types (ENH: Improve typing for decorators that change signature #80)

The more general decorator allow_kwargs is replaced with the more specific allow_only_kwargs to reduce complexity when developing (i.e., I do not have to think about any mixed-case scenarios) and to throw errors more quickly and with more informative error messages. When working with functions internally, we either want to call them with keyword-only arguments (in which case I do not need the possibility to call them without keyword arguments) or with positional arguments. In the latter case we do not gain anything by implementing an allow_only_args decorator, which is why we stick with the more general one here.

[hmgaudecker]

Looks great! Just a few renaming suggestions and general style of tests.

[hmgaudecker]

Suggested change

	@partial(vmap_1d, variables=["index"], apply_allow_kwargs=False)
	@partial(vmap_1d, variables=["index"], callable_with="kwargs")

Argument will definitely have to be renamed in vmap_1d. This is just a suggestion -- instead of Boolean with an awkward name, use Literal that could be args or kwargs. But there could be better ideas, I did not think about it for very long.

[hmgaudecker]

In simulate.py, is there a reason for having indices and index differently in the public and private version of retrieve_non_sparse_choices? If not, this might be a good occasion to harmonise them.

[timmens]

The reason was that the private function was written to work on a single index and then vmapped. Noticing that I only have to vmap jnp.unravel_index I rewrote this part for clarity.