Add `Blockwise` `Op` #1215

purna135 · 2022-09-26T16:00:42Z

This PR builds off of #757 and closes #695.

To #757 it adds:

get_output_info(), which is the same as Elemwise get_output_info(), to make all inputs of the same dimension.
derive DimShuffle's gufunc signature
reduce the broadcasted dimensions of inputs after the grad is computed

Differences with #757:

instead of using the dimensions from the start for computing the curr_static_shape of core_inp_grads use the dimensions from the end.
an extra check before calling perform() of DimShuffle (which can be removed later)

codecov · 2022-09-26T18:49:08Z

Codecov Report

Merging #1215 (6cda5c3) into main (462d8d5) will increase coverage by 4.14%.
The diff coverage is 86.00%.

❗ Current head 6cda5c3 differs from pull request most recent head c7b0d10. Consider uploading reports for the commit c7b0d10 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1215      +/-   ##
==========================================
+ Coverage   75.02%   79.16%   +4.14%     
==========================================
  Files         194      174      -20     
  Lines       50099    48677    -1422     
  Branches    12096    10359    -1737     
==========================================
+ Hits        37586    38536     +950     
+ Misses      10189     7640    -2549     
- Partials     2324     2501     +177

Impacted Files	Coverage Δ
aesara/tensor/blockwise.py	`85.81% <85.81%> (ø)`
aesara/tensor/math.py	`90.05% <100.00%> (-0.62%)`	⬇️

... and 122 files with indirect coverage changes

brandonwillard · 2022-09-26T20:00:37Z

Don't forget to rebase onto upstream/main (or whatever the name of your remote for this repository is); that should remove the merge commit in this branch.

brandonwillard

This looks great! The next step involves extending the number of gufunc_sigs we specify and adding the associated tests.

The big, open question is whether or not we can replace Elemwise with this new Op. When we demonstrate that this Op can at least handle all the standard Elemwise cases, then we'll start exploring this question further, though. In other words, we don't want to start considering all the other changes (e.g. Blockwise.c_code, Numba/JAX transpilations, etc.) until we've demonstrated good test coverage (both Elemwise/scalar broadcasting cases and otherwise).

brandonwillard · 2022-09-26T20:02:16Z

tests/tensor/test_blockwise.py

+ x = Blockwise(op)(*args)
+ x_fn = aesara.function(args, x)
+
+ x_fn(*arg_vals)


We're going to need to assert something about this output.

brandonwillard · 2022-09-26T20:03:09Z

aesara/tensor/math.py

+ gufunc_sig = ((("m", "n"), ("n", "p")), (("m", "p"),))
+
+ __props__ = ("gufunc_sig",)


FYI: We'll need to create these kinds of signatures for every applicable Op.

tests/tensor/test_blockwise.py

aesara/tensor/blockwise.py

purna135 · 2022-10-11T19:48:00Z

What should be the signature for Subtensor Op and Shape Op ?

brandonwillard · 2022-10-11T21:12:37Z

What should be the signature for Subtensor Op and Shape Op ?

If you're talking about constructing symbolic graphs, the signatures are ultimately determined by their Op.make_node implementations.

purna135 · 2022-10-11T21:15:32Z

If you're talking about constructing symbolic graphs, the signatures are ultimately determined by their Op.make_node implementations.

Yes, got it now

purna135 · 2022-10-13T08:05:02Z

Hello, @brandonwillard.
I'm having some DimShuffle related problems that I can't figure out.
Could you please take a look and assist in determining which piece of logic is causing this error?

You can reproduce the error using the following command.
pytest tests/tensor/test_blockwise.py::test_blockwise_solve_grad[a_shape0-b_shape0]

brandonwillard · 2022-10-16T00:47:57Z

Hello, @brandonwillard. I'm having some DimShuffle related problems that I can't figure out. Could you please take a look and assist in determining which piece of logic is causing this error?

You can reproduce the error using the following command. pytest tests/tensor/test_blockwise.py::test_blockwise_solve_grad[a_shape0-b_shape0]

It looks like SolveBase.L_op is producing tensors with an extra dimension that Solve2 can't handle.

I'm guessing Solve2 was meant to serve as a specialization of Solve's more general (matrix x matrix) -> matrix signature, but its inherited L_op probably doesn't match the signature change.

Regardless, we shouldn't need new Ops for that; instead, a helper function like aesara.tensor.slinalg.solve can be used to project the inputs and outputs to and from Solve's signature's space.

brandonwillard

I've added comments for some of the changes we made locally during the meeting.

brandonwillard · 2023-03-17T21:51:13Z

aesara/tensor/basic.py

+ ),
+ (("n", "m"),),
+ )
+ __props__ = ("dtype", "gufunc_sig")


Suggested change

__props__ = ("dtype", "gufunc_sig")

__props__ = ("dtype",)

brandonwillard · 2023-03-17T21:55:25Z

aesara/tensor/basic.py

@@ -3502,7 +3517,8 @@ class AllocDiag(Op):
 It does the inverse of `ExtractDiag`.
 """

- __props__ = ("offset", "axis1", "axis2")
+ gufunc_sig = (((),), (("m", "m"),))
+ __props__ = ("offset", "axis1", "axis2", "gufunc_sig")


Suggested change

__props__ = ("offset", "axis1", "axis2", "gufunc_sig")

__props__ = ("offset", "axis1", "axis2",)

brandonwillard · 2023-03-17T21:56:42Z

aesara/tensor/blockwise.py

+ return Apply(self, list(inputs), outputs)
+
+ def __str__(self):
+ return f"{type(self).__name__}{{op={self.op}}}"


Suggested change

return f"{type(self).__name__}{{op={self.op}}}"

return f"{type(self).__name__}{{{self.op}, {self.signature}}}"

brandonwillard · 2023-03-17T21:57:48Z

aesara/tensor/blockwise.py

+ # The gradient contains a constant
+ # res = aesara.tensor.basic.constant(
+ # np.asarray(var.data), dtype=var.type.dtype
+ # )
+ res = var
+
+ # TODO FIXME: Use dimensions of relevant/appropriate inputs.
+ # What exactly are those in this case?
+ nd = inputs[0].type.ndim
+
+ return atleast_Nd(res, n=nd)


Suggested change

# The gradient contains a constant

# res = aesara.tensor.basic.constant(

# np.asarray(var.data), dtype=var.type.dtype

# )

res = var

# TODO FIXME: Use dimensions of relevant/appropriate inputs.

# What exactly are those in this case?

nd = inputs[0].type.ndim

return atleast_Nd(res, n=nd)

return var

brandonwillard · 2023-03-17T21:58:26Z

aesara/tensor/slinalg.py

-
- __props__ = ("lower", "destructive", "on_error")
+ gufunc_sig = ((("m", "m"),), (("m", "m"),))
+ __props__ = ("lower", "destructive", "on_error", "gufunc_sig")


Suggested change

__props__ = ("lower", "destructive", "on_error", "gufunc_sig")

__props__ = ("lower", "destructive", "on_error",)

brandonwillard · 2023-03-17T22:00:02Z

tests/tensor/test_blockwise.py

+ from aesara.tensor.basic import Tri
+
+ blk_op = Blockwise(op=Tri(dtype="float64"), signature=(((), (), ()), (("n", "m"),)))


Suggested change

from aesara.tensor.basic import Tri

blk_op = Blockwise(op=Tri(dtype="float64"), signature=(((), (), ()), (("n", "m"),)))

blk_op = Blockwise(op=Tri(dtype="float64"))

brandonwillard · 2023-03-17T22:00:36Z

tests/tensor/test_blockwise.py

+ blk_op = Blockwise(op=Tri(dtype="float64"), signature=(((), (), ()), (("n", "m"),)))
+ out_dtype, output_shapes, inputs = blk_op.get_output_info(a, b, c)
+
+ assert out_dtype == ["float64"]


We need to assert something about output_shapes (i.e. make sure they're correct in some way).

Inspired by: aesara-devs/aesara#1215 Co-authored-by: Brandon T. Willard <[email protected]> Co-authored-by: Purna Chandra Mansingh <[email protected]> Co-authored-by: Sayam Kumar <[email protected]> Co-authored-by: Kaustubh <[email protected]>

purna135 closed this Sep 26, 2022

purna135 reopened this Sep 26, 2022

purna135 changed the title ~~Update L_op of Blockwise~~ Add Blockwise op Sep 26, 2022

brandonwillard added enhancement New feature or request important NumPy compatibility Op implementation Involves the implementation of an Op labels Sep 26, 2022

brandonwillard changed the title ~~Add Blockwise op~~ Add Blockwise Op Sep 26, 2022

brandonwillard mentioned this pull request Sep 26, 2022

Add Blockwise Op #757

Closed

brandonwillard reviewed Sep 26, 2022

View reviewed changes

purna135 force-pushed the add_blockwise branch from a57528c to f770ada Compare September 28, 2022 16:45

brandonwillard mentioned this pull request Sep 29, 2022

Update L_op of Blockwise brandonwillard/aesara#3

Closed

brandonwillard force-pushed the add_blockwise branch from f770ada to 6cda5c3 Compare October 6, 2022 19:19

brandonwillard self-assigned this Oct 6, 2022

purna135 commented Oct 11, 2022

View reviewed changes

aesara/tensor/blockwise.py Show resolved Hide resolved

purna135 force-pushed the add_blockwise branch 7 times, most recently from 0792e8a to fdb3045 Compare November 4, 2022 20:54

brandonwillard force-pushed the add_blockwise branch from 877d04d to c9ad602 Compare November 11, 2022 00:53

brandonwillard and others added 6 commits March 17, 2023 12:16

Fix AllocDiag and Tri gufunc signatures

be1b794

Fixed output dtype

f0eb8f5

Fix a core inputs computation bug and do some refactoring

de3c1ea

add more tests to test_infer_shape_to_gufunc_sig

45f5eeb

fix infer_shape_to_gufunc_sig

7f1b99d

add test for Blockwise SolveTriangular

c7b0d10

brandonwillard force-pushed the add_blockwise branch from 3ed3497 to c7b0d10 Compare March 17, 2023 17:16

brandonwillard reviewed Mar 17, 2023

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `Blockwise` `Op` #1215

Add `Blockwise` `Op` #1215

purna135 commented Sep 26, 2022 •

edited by brandonwillard

codecov bot commented Sep 26, 2022 •

edited

brandonwillard commented Sep 26, 2022 •

edited

brandonwillard left a comment

brandonwillard Sep 26, 2022

brandonwillard Sep 26, 2022

purna135 commented Oct 11, 2022

brandonwillard commented Oct 11, 2022

purna135 commented Oct 11, 2022

purna135 commented Oct 13, 2022

brandonwillard commented Oct 16, 2022

brandonwillard left a comment

brandonwillard Mar 17, 2023

brandonwillard Mar 17, 2023

brandonwillard Mar 17, 2023

brandonwillard Mar 17, 2023

brandonwillard Mar 17, 2023

brandonwillard Mar 17, 2023

brandonwillard Mar 17, 2023

		gufunc_sig = ((("m", "n"), ("n", "p")), (("m", "p"),))

		__props__ = ("gufunc_sig",)

	__props__ = ("offset", "axis1", "axis2", "gufunc_sig")
	__props__ = ("offset", "axis1", "axis2",)

	return f"{type(self).__name__}{{op={self.op}}}"
	return f"{type(self).__name__}{{{self.op}, {self.signature}}}"

	__props__ = ("lower", "destructive", "on_error", "gufunc_sig")
	__props__ = ("lower", "destructive", "on_error",)

		from aesara.tensor.basic import Tri

		blk_op = Blockwise(op=Tri(dtype="float64"), signature=(((), (), ()), (("n", "m"),)))

Add Blockwise Op #1215

Are you sure you want to change the base?

Add Blockwise Op #1215

Conversation

purna135 commented Sep 26, 2022 • edited by brandonwillard

codecov bot commented Sep 26, 2022 • edited

Codecov Report

brandonwillard commented Sep 26, 2022 • edited

brandonwillard left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

purna135 commented Oct 11, 2022

brandonwillard commented Oct 11, 2022

purna135 commented Oct 11, 2022

purna135 commented Oct 13, 2022

brandonwillard commented Oct 16, 2022

brandonwillard left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Add `Blockwise` `Op` #1215

Add `Blockwise` `Op` #1215

purna135 commented Sep 26, 2022 •

edited by brandonwillard

codecov bot commented Sep 26, 2022 •

edited

brandonwillard commented Sep 26, 2022 •

edited