Add ability to set tf.function args for scipy optimizer #2064

khurram-ghani · 2023-04-21T13:23:29Z

PR type: enhancement

Related issue(s)/PRs: None

Summary

Proposed changes

The main aim of this PR is to add XLA compilation support for the Scipy optimizer. This can be achieved by passing jit_compile=True argument to tf.function.
However the implemented mechanism is more general, allowing any argument to be passed to tf.function call in Scipy.minimize.
Note the new argument added to Scipy.minimize is a normal dictionary and not keyword-args. There is already a kwargs argument for scipy optimizer itself.

What alternatives have you considered?

Considered changing the existing compile: bool argument to Scipy.minimize() to an enum; with options for "no-compilation", "default-compilation" and "XLA-compilation" -- or something similar. However, decided that it is better to add a more general purpose mechanism allowing setting of any tf.function argument. There are other existing arguments (and there might be others in the future) that some users may want to set.

Minimal working example

# Enable XLA compilation
opt = gpflow.optimizers.Scipy()
result = opt.minimize(..., tffun_args={'jit_compile': True})

Release notes

Fully backwards compatible: yes

PR checklist

uri-granta

LGTM! Just one comment.

uri-granta · 2023-04-24T12:16:23Z

gpflow/optimizers/scipy.py

        step_callback: Optional[StepCallback] = None,
        compile: bool = True,
        allow_unused_variables: bool = False,
+        tffun_args: Dict[str, Any] = {},


More general to allow any mapping, not just dicts.

Generally best to avoid mutable default arguments, since if you mutate them in (or outside) the function they end up being changed for all subsequent calls. For dicts it's conventional to use None (since there isn't a built in immutable dict structure):

Suggested change

tffun_args: Dict[str, Any] = {},

tffun_args: Optional[Mapping[str, Any]] = None,

....

if ttfun_args is None:

ttfun_args = {}

I reasoned that tffun_args wasn't ever going to be mutated; but I agree it is better to use None. I have made both those suggested improvements.

uri-granta · 2023-04-24T12:18:48Z

gpflow/optimizers/scipy.py

        variables: Sequence[tf.Variable],
        compile: bool = True,
        allow_unused_variables: bool = False,
+        tffun_args: Dict[str, Any] = {},


uri-granta · 2023-04-24T12:20:07Z

tests/gpflow/optimizers/test_scipy.py

+    "tffun_args",
+    [{}, dict(jit_compile=True), dict(jit_compile=False, other_arg="dummy")],
+)
+def test_scipy__tffun_args(


awav · 2023-04-24T13:06:23Z

Are we sure that tffun_args is a good name?
Which use case do you want to cover?

Note the new argument added to Scipy.minimize is a normal dictionary and not keyword-args. There is already a kwargs argument for scipy optimizer itself.

Yes. Wouldn't adding a dictionary with arguments to tensorflow tf.function be even more confusing?

khurram-ghani · 2023-04-24T14:00:55Z

Are we sure that tffun_args is a good name?

I am not. Suggestions welcome :-)

Which use case do you want to cover?

The current version of tf.function supports arguments such as reduce_retracing, autograph, experimental_autograph_options that users may want to set. Also jit_compile used to be called experimental_compile previously. So argument names may change in the future and I thought it might be better to not burn them into gpflow.

However if we want to keep it simple, we could just add direct support for XLA compilation by, e.g., modifying the compile argument.

Note the new argument added to Scipy.minimize is a normal dictionary and not keyword-args. There is already a kwargs argument for scipy optimizer itself.

Yes. Wouldn't adding a dictionary with arguments to tensorflow tf.function be even more confusing?

It is a bit messy/confusing to have two different dictionaries. I am not sure how to add a generic mechanism in a backwards compatible way. We could combine the two dictionaries into one (but if we rename the existing scipy_kwargs, it won't be backwards compatible), but that could be confusing itself as the user can't be completely sure where the arguments are going.

uri-granta · 2023-04-25T10:32:00Z

Are we sure that tffun_args is a good name?

Wouldn't adding a dictionary with arguments to tensorflow tf.function be even more confusing?

Personally, I think the current approach is the best one (though the docstring for ttfun_args should be expanded to include an example of how to use it for XLA).

A generic thin wrapper is easier to maintain and more future proof than an adhoc XLA option.
We already pass through arbitrary arguments to the scipy optimizer.
The name ttfun_args is very much in the style of scipy_kwargs.
The options parameter often used in scipy_kwargs is already an explicit dictionary.

I therefore don't think tffun_args being a dictionary is an issue, especially if we give an example in the docstring. Though we should perhaps use the dict syntax to make it even clearer, as in:

minimize(
        m3.training_loss,
        variables=m3.trainable_variables,
        options=dict(maxiter=50),
        compile=True,
        tffun_args=dict(jit_compile=True),
    )

khurram-ghani · 2023-04-25T10:47:49Z

I have added an example to the docstring.

gpflow/optimizers/scipy.py

codecov · 2023-04-25T11:59:02Z

Codecov Report

Patch coverage: 100.00% and project coverage change: +0.01 🎉

Comparison is base (efade00) 98.16% compared to head (eac1f8f) 98.18%.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #2064      +/-   ##
===========================================
+ Coverage    98.16%   98.18%   +0.01%     
===========================================
  Files           97       97              
  Lines         5458     5462       +4     
===========================================
+ Hits          5358     5363       +5     
+ Misses         100       99       -1

Impacted Files	Coverage Δ
gpflow/optimizers/scipy.py	`98.16% <100.00%> (+1.02%)`	⬆️

☔ View full report in Codecov by Sentry.
📢 Do you have feedback about the report comment? Let us know in this issue.

…n_args

Helps with codecov.

…n_args

khurram-ghani added 5 commits April 20, 2023 17:56

Add ability to set tf.function args for scipy opt

37b51fe

Update release notes

ad1469e

Select XLA option depending on tf version

a54d4d7

Fix minor format

2f53378

Add check for tffun_args without compile

0b894cb

uri-granta self-requested a review April 24, 2023 12:08

uri-granta approved these changes Apr 24, 2023

View reviewed changes

awav self-requested a review April 24, 2023 13:03

Improve tffun_args type and default

ca80538

Add jit_compile example in docs

9592cc1

khurram-ghani commented Apr 25, 2023

View reviewed changes

gpflow/optimizers/scipy.py Outdated Show resolved Hide resolved

khurram-ghani added 4 commits April 25, 2023 13:01

Merge remote-tracking branch 'origin/develop' into khurram/scipy_tffu…

3adabdf

…n_args

Make new arg non-optional for internal method

10de500

Helps with codecov.

Merge remote-tracking branch 'origin/develop' into khurram/scipy_tffu…

fb3af46

…n_args

Rename new arg

eac1f8f

khurram-ghani merged commit 3cfa81b into develop May 3, 2023

khurram-ghani deleted the khurram/scipy_tffun_args branch May 3, 2023 08:26

khurram-ghani mentioned this pull request May 3, 2023

Release 2.8.0 #2068

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add ability to set tf.function args for scipy optimizer #2064

Add ability to set tf.function args for scipy optimizer #2064

Uh oh!

khurram-ghani commented Apr 21, 2023

Uh oh!

uri-granta left a comment

Uh oh!

uri-granta Apr 24, 2023

Uh oh!

khurram-ghani Apr 24, 2023

Uh oh!

uri-granta Apr 24, 2023

Uh oh!

uri-granta Apr 24, 2023

Uh oh!

awav commented Apr 24, 2023 •

edited

Loading

Uh oh!

khurram-ghani commented Apr 24, 2023

Uh oh!

uri-granta commented Apr 25, 2023

Uh oh!

khurram-ghani commented Apr 25, 2023

Uh oh!

Uh oh!

codecov bot commented Apr 25, 2023 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

-        tffun_args: Dict[str, Any] = {},
+        tffun_args: Optional[Mapping[str, Any]] = None,
+    ....
+        if ttfun_args is None:
+            ttfun_args = {}

Add ability to set tf.function args for scipy optimizer #2064

Add ability to set tf.function args for scipy optimizer #2064

Uh oh!

Conversation

khurram-ghani commented Apr 21, 2023

Summary

Minimal working example

Release notes

PR checklist

Uh oh!

uri-granta left a comment

Choose a reason for hiding this comment

Uh oh!

uri-granta Apr 24, 2023

Choose a reason for hiding this comment

Uh oh!

khurram-ghani Apr 24, 2023

Choose a reason for hiding this comment

Uh oh!

uri-granta Apr 24, 2023

Choose a reason for hiding this comment

Uh oh!

uri-granta Apr 24, 2023

Choose a reason for hiding this comment

Uh oh!

awav commented Apr 24, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

khurram-ghani commented Apr 24, 2023

Uh oh!

uri-granta commented Apr 25, 2023

Uh oh!

khurram-ghani commented Apr 25, 2023

Uh oh!

Uh oh!

codecov bot commented Apr 25, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

awav commented Apr 24, 2023 •

edited

Loading

codecov bot commented Apr 25, 2023 •

edited

Loading