feat: Add overload signatures to module.apply

[Daraan]

What does this PR do?

The apply method returns either the output or a tuple[output, dict] depending on mutable (in core.apply).

flax/flax/core/scope.py

Lines 1080 to 1083 in 7d4a288

	if mutable is not False:
	return y, root.mutable_variables()
	else:
	return y

The current type hints are were general in the return type -> Any | tuple[Any, FrozenVariableDict | dict[str, Any]]: which can be confusing lead to type-checking errors. A related SO question how to type Module.apply correctly.

This PR adds two typing.overload signatures to Module.apply to differentiate these cases - when to expect a tuple and when not

See also #2086 as a related discussion.

Resolves:
#4779

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other
  checks if that's the case).
• This change is discussed in a Github issue/
  discussion (please add a
  link).

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.