Jax ml operator fix #4041
Jax ml operator fix #4041
Conversation
|
|
|
@pbrubeck @dham Firedrake pip installs PR depends on the approval of UFL PR 348 and this PR. |
|
The failure in the docs building is because of a recent Sphinx update. |
| @@ -77,6 +77,6 @@ def J(f): | |||
| c = Control(f) | |||
| Jhat = ReducedFunctional(J(f), c) | |||
|
|
|||
| f_opt = minimize(Jhat, tol=1e-6, method="BFGS") | |||
| f_opt = minimize(Jhat, tol=1e-4, method="BFGS") | |||
There was a problem hiding this comment.
This tolerance modification reduces the runtime test by almost 50%. Most important is that the following assert is still satisfied.
firedrake/ml/jax/ml_operator.py
Outdated
| @@ -39,36 +39,41 @@ def __init__( | |||
| *operands: Union[ufl.core.expr.Expr, ufl.form.BaseForm], | |||
| function_space: WithGeometryBase, | |||
| derivatives: Optional[tuple] = None, | |||
| argument_slots: Optional[tuple[Union[ufl.coefficient.BaseCoefficient, ufl.argument.BaseArgument]]], | |||
| argument_slots: Optional[tuple[Union[ufl.coefficient.BaseCoefficient, ufl.argument.BaseArgument]]] = (), | |||
There was a problem hiding this comment.
If argument slots are not provided, then they will be generated.
Having an empty tuple seems a strange default considering. Could/should it be None instead?
There was a problem hiding this comment.
Actually if it is Optional then that's what is has to be (https://docs.python.org/3/library/typing.html#typing.Optional)
There was a problem hiding this comment.
Optional[tuple[Union[ufl.coefficient.BaseCoefficient, ufl.argument.BaseArgument]]] = None leads to the error:
> argument_slots = tuple(map(as_ufl, argument_slots))
E TypeError: 'NoneType' object is not iterable
../ufl/ufl/core/base_form_operator.py:49: TypeError
There was a problem hiding this comment.
The way to define argument_slots argument is comming from https://github.com/firedrakeproject/ufl/blob/master/ufl/core/external_operator.py#L28
There was a problem hiding this comment.
We can improve this by writing
argument_slots: tuple[Union[ufl.coefficient.BaseCoefficient, ufl.argument.BaseArgument]] | None = None,
Then
if argument_slots is None:
argument_slots = ()
There was a problem hiding this comment.
I think it would be better to just have
argument_slots: tuple[Union[ufl.coefficient.BaseCoefficient, ufl.argument.BaseArgument]] = (),(i.e. no Optional) and to remove
If argument slots are not provided, then they will be generated.
from the docstring.
That is at least consistent with the UFL object.
Description
Depends on firedrakeproject/ufl#60
Fixed an issue in
firedrake.ml.jax.ml_operatorwhereargument_slotswas incorrectly specified. I wrote it as optional:Clarified that if
argument_slotsis not provided,ML_Operatorwill automatically write it.Test the results involving the Neural operators are in right function space.
Test jax and pytorch operators in Firedrake CI.