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.

tests/tensor/test_blockwise.py

brandonwillard · 2023-03-10T21:03:01Z

aesara/tensor/blockwise.py

+    """
+    op = node.op
+    in_shapes = tuple(
+        tuple(lscalar(f"i{s}") for s in range(inp.type.ndim)) for inp in node.inputs


Looks like we need to change the labels to reflect the inputs. For example:

Suggested change

tuple(lscalar(f"i{s}") for s in range(inp.type.ndim)) for inp in node.inputs

tuple(lscalar(f"i{n}{s}") for s in range(inp.type.ndim)) for n, inp in enumerate(node.inputs)

Co-authored-by: Brandon T. Willard <[email protected]> Co-authored-by: Sayam Kumar <[email protected]> Co-authored-by: Kaustubh <[email protected]>

… gufunc

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


-    __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).

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 force-pushed the add_blockwise branch from 9f973b9 to 4a926fa Compare January 21, 2023 22:27

brandonwillard force-pushed the add_blockwise branch from f045067 to a5de97c Compare February 2, 2023 22:57

brandonwillard force-pushed the add_blockwise branch from b2b2604 to c2470b5 Compare February 19, 2023 19:00

brandonwillard reviewed Mar 10, 2023

View reviewed changes

tests/tensor/test_blockwise.py Show resolved Hide resolved

brandonwillard reviewed Mar 10, 2023

View reviewed changes

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

Extract check method from InferShapeTester

b6c0e0e

Add a Blockwise Op

5037958

Co-authored-by: Brandon T. Willard <[email protected]> Co-authored-by: Sayam Kumar <[email protected]> Co-authored-by: Kaustubh <[email protected]>

simply printing of Blockwise Op

f8af028

refactor grad_signature

2f0576d

use Blockwise instead of Elemwise

2c82292

Fix typing issues in Blockwise._bgrad

0bb4feb

Convert inputs to ndarrays in Blockwise.perform

f67c769

Add a Blockwise Op

90f4791

fixed py_func result

d24f619

fixed parameters to atleast_Nd()

daf2ea2

added gufunc_sig to Tri and AllocDiag

33ead95

added gufunc_sig to Cholesky

fdf5af8

added test for Blockwise Cholesky

0072f2d

manage the Ops which support nd inputs

ab66791

Use dispatch for gufunc signature and (partially) implement Subtensor…

b77d50b

… gufunc

Add a infer_shape_to_gufunc_sig function

76dc6b4

Fix flat_out_shapes

e1c8892

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

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

		__props__ = ("gufunc_sig",)

	tuple(lscalar(f"i{s}") for s in range(inp.type.ndim)) for inp in node.inputs
	tuple(lscalar(f"i{n}{s}") for s in range(inp.type.ndim)) for n, inp in enumerate(node.inputs)

	__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"),)))

Uh oh!

Add Blockwise Op #1215

Are you sure you want to change the base?

Add Blockwise Op #1215

Uh oh!

Conversation

purna135 commented Sep 26, 2022 • edited by brandonwillard Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Sep 26, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

brandonwillard commented Sep 26, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

brandonwillard left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

purna135 commented Oct 11, 2022

Uh oh!

brandonwillard commented Oct 11, 2022

Uh oh!

purna135 commented Oct 11, 2022

Uh oh!

purna135 commented Oct 13, 2022

Uh oh!

brandonwillard commented Oct 16, 2022

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

brandonwillard left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Add `Blockwise` `Op` #1215

Add `Blockwise` `Op` #1215

purna135 commented Sep 26, 2022 •

edited by brandonwillard

Loading

codecov bot commented Sep 26, 2022 •

edited

Loading

brandonwillard commented Sep 26, 2022 •

edited

Loading