google / google/flax

Start using positional-only parameters in Linen

Open
#2,049 2 comments 0 reactions 0 assignees View on GitHub
Priority: P2 - no schedule
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

Some functions in our public API use positional arguments and keyword arguments in a way that is somewhat error prone. Two examples:

```py
# Example 1 (module.py)
def variable(self, col: str, name: str,
init_fn: Optional[Callable[..., Any]] = None,
*init_args) -> Variable:
...
```

Calling this function with `variable('a', 'b', 'arg1', init_fn=fn)` will give two vaues to `init_fn`, which is wrong.

```
# Example 2 (module.py)
def init(self,
rngs: Union[PRNGKey, RNGSequences],
*args,
method: Optional[Callable[..., Any]] = None,
mutable: CollectionFilter = DenyList('intermediates'),
**kwargs) -> FrozenVariableDict:
```

If users pass `rng=` as a kwarg it will break while they probably wanted to forward a rng keyword arg to their `__call__` function.

We could resolve this by start using Python's 3.8 "positonal-only arguments: https://docs.python.org/3/whatsnew/3.8.html#positional-only-parameters

So that will look as follows:

```
def variable(self, col: str, name: str,
init_fn: Optional[Callable[..., Any]] = None,
/,
*init_args) -> Variable:
```

Now `variable('a', 'b', 'arg1', init_fn=fn)` will give error `TypeError: variable() got some positional-only arguments passed as keyword arguments: 'init_fn'`.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.