Removes any kwargs not supported by `module_or_fn` from `all_kwargs_dict`.

  A new dict is return with shallow copies of keys & values from
  `all_kwargs_dict`, as long as the key is accepted by module_or_fn. The
  returned dict can then be used to connect `module_or_fn` (along with some
  other inputs, ie non-keyword arguments, in general).

  `snt.supports_kwargs` is used to tell whether a given kwarg is supported. Note
  that this method may give false negatives, which would lead to extraneous
  removals in the result of this function. Please read the docstring for
  `snt.supports_kwargs` for details, and manually inspect the results from this
  function if in doubt.

  Args:
    module_or_fn: some callable which can be interrogated by
      `snt.supports_kwargs`. Generally a Sonnet module or a method (wrapped in
      `@reuse_variables`) of a Sonnet module.
    all_kwargs_dict: a dict containing strings as keys, or None.

  Raises:
    ValueError: if `all_kwargs_dict` is not a dict.

  Returns:
    A dict containing some subset of the keys and values in `all_kwargs_dict`.
    This subset may be empty. If `all_kwargs_dict` is None, this will be an
    empty dict.

def remove_unsupported_kwargs(module_or_fn, all_kwargs_dict):
  """Removes any kwargs not supported by `module_or_fn` from `all_kwargs_dict`.

  A new dict is return with shallow copies of keys & values from
  `all_kwargs_dict`, as long as the key is accepted by module_or_fn. The
  returned dict can then be used to connect `module_or_fn` (along with some
  other inputs, ie non-keyword arguments, in general).

  `snt.supports_kwargs` is used to tell whether a given kwarg is supported. Note
  that this method may give false negatives, which would lead to extraneous
  removals in the result of this function. Please read the docstring for
  `snt.supports_kwargs` for details, and manually inspect the results from this
  function if in doubt.

  Args:
    module_or_fn: some callable which can be interrogated by
      `snt.supports_kwargs`. Generally a Sonnet module or a method (wrapped in
      `@reuse_variables`) of a Sonnet module.
    all_kwargs_dict: a dict containing strings as keys, or None.

  Raises:
    ValueError: if `all_kwargs_dict` is not a dict.

  Returns:
    A dict containing some subset of the keys and values in `all_kwargs_dict`.
    This subset may be empty. If `all_kwargs_dict` is None, this will be an
    empty dict.
  """
  if all_kwargs_dict is None:
    all_kwargs_dict = {}
  if not isinstance(all_kwargs_dict, dict):
    raise ValueError("all_kwargs_dict must be a dict with string keys.")
  return {
      kwarg: value for kwarg, value in all_kwargs_dict.items()
      if supports_kwargs(module_or_fn, kwarg) != NOT_SUPPORTED
  }