Support expanding some functions #125

ZibingZhang · 2022-11-19T23:30:09Z

Overview

Adds a NodeTransformer that expands functions.

Details

Adds a expand_functions keyword argument onto get_latex which allows the user to specify the names of the methods they want to have expanded.
Introduces a class FunctionExpander(ast.NodeExpander) which replaces ast.Call with the equivalent expanded ast.AST.
Only expanded method to be supported is hypot(x, y) -> \sqrt{x^2 + y^2}.
Supports nested hypot calls, e.g. hypot(hypot(a, b), x, y).

Thoughts

As discussed in the issue, there may be multiple ways to expand a function. We could expose various expansion methods, and accept them in the form expanded_functions={"method1": Method1Expander1, "method2": None} where None defaults to some predetermined expansion. This is a bit gross, but we could probably make this a bit nicer with the config refactor. This would also allow users to define their own expansion methods.

There isn't a way to prevent a stack overflow right now due to f(x) -> g(x) and g(x) -> f(x). My initial thought would be to add a function_expansion_depth=1 as a keyword argument, and it would be implemented like so:

Initialize current_depth to 0 for each instance of FunctionExpander.
In visit_Call:
- If current_depth > function_expansion_depth, then return early.
- If any expansions are applied, increment current_depth by 1 (I'd err towards forcing users to specify if they want to try for deeper level of replacements).
  - After mapping over child nodes decrement current_depth by 1.

References

#65

Blocked by

Potentially blocked by #103.

We could merge this in before, and then figure out a config format in the other PR, or vice versa.

google-cla · 2022-11-19T23:30:12Z

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.

odashi

This is very good, and sorry for putting many comments (most of them are trivial I think)

src/integration_tests/regression_test.py

src/latexify/frontend.py

src/latexify/transformers/function_expander.py

odashi · 2022-11-20T09:07:19Z

src/integration_tests/regression_test.py

+        r"\sqrt{a^{{2}} + b^{{2}}}^{{2}} + "
+        r"x^{{2}} + y^{{2}}}"


nit: Maybe it'd be good to apply some symbolic optimizations since nested hypot can be removed: $\sqrt{{\sqrt{x^{2} + y^{2}}}^{2} + z^{2}} = \sqrt{x^{2} + y^{2} + z^{2}}$ .

This is beyond the scope of this pull request. We may need to make another issue to discuss about it.

I would imagine that this would be a separate transformer that takes place before function expansion

Yeah, and it would become a more complex optimizer. We could borrow some knowledge from other symbolic math tools as well.

src/latexify/codegen/latex_test.py

src/latexify/test_stubs.py

src/latexify/transformers/function_expander_test.py

odashi · 2022-11-20T09:51:49Z

There isn't a way to prevent a stack overflow right now due to f(x) -> g(x) and g(x) -> f(x)

Nice insight. For now we can introduce functions that don't employ recursive definitions, but we need to keep in mind that this could become a problem in the future. It is better to add some TODOs to function_expander.py.

ZibingZhang · 2022-11-20T11:45:36Z

src/latexify/transformers/function_expander.py

+
+    args_reduced = functools.reduce(lambda a, b: ast.BinOp(a, ast.Add(), b), args)
+    return ast.Call(
+        func=ast.Name(id="sqrt", ctx=Load()),


Maybe in constants.py, BUILTIN_FUNCS should be dict[enum, tuple[str, str]]?

Basically what's happening here is that we're replacing hypot with sqrt, while also relying on the fact sqrt gets further preprocessed by getting replaced with (r"\sqrt{", "}").

Having an enum rather than a bunch of strings might make the code a bit more readable, as well as helping to prevent errors with typos. In the future with more function expansions and different transformers, there might be a lot of cases where we care about specific, commonly-used functions, such as sqrt.

If you agree with changing it to an enum I'll tackle that in another issue / PR.

To be clear, do you mean defining an Enum class something like:

class BuiltinFnName(str, enum.Enum): ... SQRT = "sqrt"

and use it in:

BUILTIN_FUNCS: dict[BuiltinFnName, tuple[str, str]] = { ... BuiltinFnName.SQRT: (r"\sqrt{", "}"), ... }

? If I'm wrong please fix me.

something like that yeah

It looks nicer than the current impl. I guess we finally need to adopt more sophisticated way to represent functions, but for now the idea above looks a good choice.

odashi

Sorry for bothering again! Overall I like this change, but there are some comments before merging.

A general suggestion: if we instantiate AST nodes, it is preferred to explicitly specify parameter names for every argument. AST structure is sometimes updated, and without parameter names it may lack version compatibility (not only currently supported versions, but also versions in the future, in which we couldn't expect what kind of changes will be introduced)

src/latexify/transformers/function_expander.py

src/latexify/transformers/function_expander_test.py

odashi · 2022-11-20T20:20:46Z

src/latexify/ast_utils_test.py

+        (ast.Call(ast.Name("hypot"), ast.Constant("foo")), "hypot"),
+        (
+            ast.Call(
+                ast.Attribute(ast.Name("math"), "hypot", ast.Load()),
+                ast.Constant("foo"),
+            ),
+            "hypot",
+        ),
+        (ast.Call(ast.Constant(123), ast.Constant("foo")), None),


These Calls are illegal because the args parameter should take a list of expr.

The func parameter could take any expr, but in this case the constant 123 is never populated here in the actual AST. I guess some other example is preferred, e.g., ast.Call(func=ast.Name("foo"), args=[]), which represents foo()(...)

ZibingZhang · 2022-11-20T20:31:51Z

Sorry for bothering again! Overall I like this change, but there are some comments before merging.

No worries at all! Thanks for taking the time to review :) I think this is a really cool project.

I'll take a look at this sometime later today / tomorrow. Do you have any recs for another issue for me to look at after this one? It seems like quite a few of them are blocked my issues that have existing PRs, but I'd love to work on something else if I've got free time this week

odashi · 2022-11-21T00:06:09Z

@ZibingZhang I think this PR is almost independent from others except #103 which is an essential change. If you could take your time for that PR it's welcome.

odashi · 2022-11-21T00:12:12Z

I also think adding supports for other functions in math and numpy can be processed independently (but blocked by this issue #125).

ZibingZhang · 2022-11-21T09:01:46Z

A general suggestion: if we instantiate AST nodes, it is preferred to explicitly specify parameter names for every argument.

Good point, will keep this in mind moving forward.

I also think adding supports for other functions in math and numpy can be processed independently (but blocked by this issue #125).

That's a good point. I'll start working through that backlog.

odashi

Really good change, thanks!

Impl function expander

a0da875

ZibingZhang requested a review from odashi as a code owner November 19, 2022 23:30

Zibing Zhang added 3 commits November 20, 2022 00:09

Fix error message

7c6701e

fix flake8 errors

a220960

file doc

89ad355

odashi reviewed Nov 20, 2022

View reviewed changes

Zibing Zhang added 6 commits November 20, 2022 10:34

Suggestions

644b4ec

Fix docstring

46d4a8a

Fix docstring

9eec975

rename test methods

f4dcfd0

slight rephrasing

819e8be

rm local var

0797743

ZibingZhang requested a review from odashi November 20, 2022 10:46

Zibing Zhang added 2 commits November 20, 2022 10:47

Merge branch 'master'

7fdea83

Reformat import

df5a64e

ZibingZhang commented Nov 20, 2022

View reviewed changes

odashi reviewed Nov 20, 2022

View reviewed changes

suggestions

d8cba78

ZibingZhang requested a review from odashi November 21, 2022 09:01

odashi approved these changes Nov 21, 2022

View reviewed changes

odashi merged commit 85dddda into google:main Nov 21, 2022

odashi mentioned this pull request Nov 21, 2022

Support expanding some functions #65

Closed

This was referenced Nov 21, 2022

Add BuiltinFnEnum #130

Merged

Add keyword to all AST node constructor positional arguments. #131

Merged

odashi added the feature label Nov 21, 2022

odashi added this to the v0.3 milestone Nov 21, 2022

ZibingZhang deleted the main branch November 22, 2022 14:29

odashi mentioned this pull request Nov 23, 2022

Move expand_functions to Config #133

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support expanding some functions #125

Support expanding some functions #125

ZibingZhang commented Nov 19, 2022 •

edited

Loading

google-cla bot commented Nov 19, 2022

odashi left a comment

odashi Nov 20, 2022

ZibingZhang Nov 20, 2022

odashi Nov 20, 2022

odashi commented Nov 20, 2022

ZibingZhang Nov 20, 2022

odashi Nov 20, 2022 •

edited

Loading

ZibingZhang Nov 20, 2022

odashi Nov 20, 2022

odashi left a comment

odashi Nov 20, 2022

ZibingZhang commented Nov 20, 2022

odashi commented Nov 21, 2022

odashi commented Nov 21, 2022

ZibingZhang commented Nov 21, 2022

odashi left a comment

Support expanding some functions #125

Support expanding some functions #125

Conversation

ZibingZhang commented Nov 19, 2022 • edited Loading

Overview

Details

Thoughts

References

Blocked by

google-cla bot commented Nov 19, 2022

odashi left a comment

Choose a reason for hiding this comment

odashi Nov 20, 2022

Choose a reason for hiding this comment

ZibingZhang Nov 20, 2022

Choose a reason for hiding this comment

odashi Nov 20, 2022

Choose a reason for hiding this comment

odashi commented Nov 20, 2022

ZibingZhang Nov 20, 2022

Choose a reason for hiding this comment

odashi Nov 20, 2022 • edited Loading

Choose a reason for hiding this comment

ZibingZhang Nov 20, 2022

Choose a reason for hiding this comment

odashi Nov 20, 2022

Choose a reason for hiding this comment

odashi left a comment

Choose a reason for hiding this comment

odashi Nov 20, 2022

Choose a reason for hiding this comment

ZibingZhang commented Nov 20, 2022

odashi commented Nov 21, 2022

odashi commented Nov 21, 2022

ZibingZhang commented Nov 21, 2022

odashi left a comment

Choose a reason for hiding this comment

ZibingZhang commented Nov 19, 2022 •

edited

Loading

odashi Nov 20, 2022 •

edited

Loading