Add name kwarg to Op.call #693

HarshvirSandhu · 2024-04-07T19:15:27Z

Description

Add name keyword in Op.__call__ allowing arbitrary Ops to be given names

Related Issue

Closes Add name keyword argument to Op.__call__ #685

Checklist

Included tests that prove the fix is effective or that the new feature works

Type of change

New feature / enhancement

pytensor/graph/op.py

ricardoV94 · 2024-04-12T06:28:27Z

pytensor/graph/op.py

        node = self.make_node(*inputs, **kwargs)
+        if isinstance(node.outputs, list):


This will always be a list.

More importantly I think we should only assign name to the default output (if there's a single output, that's the default).

If there's no default output we could call them "f{name}_{i}" perhaps? What do you think

If my understanding is correct

If node.outputs has a length of more than 1, then they should be called f"{name}_{i}"
Here i would be the default_output(which is specified in Multi-output Ops)

Not quite what I was thinking. Multi-output nodes can have a default output, which is the only thing users usually see. We should pass name directly to either to the default output of a multi-output node or the only output of a single-output node.

But for multi-output nodes without a default output we shouldn't. For those I suggest adding a numerical suffix, 0 for the first, 1 for the second, and so on

ricardoV94

Small tweaks/suggestion

ricardoV94

Also we need tests for multiple output with default output and not.

There are some dummy outputs we use to test basic functionality that doesn't require finding real ops.

Should be somewhere in tests/graph/....py

ricardoV94 · 2024-04-13T08:00:10Z

pytensor/graph/op.py

+            if len(node.outputs) == 1:
+                node.outputs[0].name = name
+            else:
+                for i, n in enumerate(node.outputs):
+                    n.name = f"{name}_{i}"


Suggested change

if len(node.outputs) == 1:

node.outputs[0].name = name

else:

for i, n in enumerate(node.outputs):

n.name = f"{name}_{i}"

if len(node.outputs) == 1:

node.outputs[0].name = name

elif self.default_output is not None:

node.outputs[self.default_output].name = name

else:

for i, n in enumerate(node.outputs):

n.name = f"{name}_{i}"

ricardoV94 · 2024-04-13T08:02:30Z

Also scan seems to have a special logic for the name kwarg (see failing tests). Have to check what's going on with that

HarshvirSandhu · 2024-04-13T08:11:19Z

Also scan seems to have a special logic for the name kwarg (see failing tests). Have to check what's going on with that

For scan some tests were failing because they were assigned names like None_1
I have put a condition to assign names only if name is not None

ricardoV94 · 2024-04-13T08:18:36Z

tests/graph/test_op.py

+    for i, r in enumerate(res):
+        assert r.name == f"{op_name}_{i}"
+
+    z = pt.add(x, y, name=op_name)


Don't use pt.add. Use a single output dummy Op

ricardoV94 · 2024-04-13T08:19:30Z

tests/graph/test_op.py

+            return Apply(self, list(inputs), outputs)
+
+        def perform(self, node, inputs, outputs):
+            outputs[0] = pt.matrix()


This is wrong but not needed. You can replace contents of perform with raise NotImplementedError

tests/graph/test_op.py

ricardoV94 · 2024-04-13T08:21:14Z

tests/graph/test_op.py

+
+
+def test_op_name():
+    x = pt.vector("x")


Use dummy test variables instead of pt.vector (the reason is this tests are in core abstract functionality). tensors ops and variables are specific implementations of these abstract objects

ricardoV94 · 2024-04-13T08:23:36Z

pytensor/graph/op.py

@@ -289,7 +289,14 @@ def __call__(self, *inputs: Any, **kwargs) -> Variable | list[Variable]:

        """
        return_list = kwargs.pop("return_list", False)
+        name = kwargs.pop("name", None)


Instead of popping we can make it an explicit optional kwarg in the call signature (so it's actually discoverable). Same for return_list. No idea why they went for this implicit approach

ricardoV94 · 2024-04-15T09:18:11Z

tests/graph/test_op.py

+    class DummyType(Type):
+        def filter(self, data):
+            return data
+
+        def __eq__(self, other):
+            return isinstance(other, DummyType)
+
+        def __hash__(self):
+            return hash(DummyType)
+
+        def __repr__(self):
+            return "DummyType()"


Do we need a new dummy type? Can we reuse one from the existing ones for testing?

The existing dummy type takes an argument called: thingy

I'm not sure if that should be reused (I assume it has a specific use)

Should be fine, it's a test Type, it takes an argument for other purposes but that shouldn't be problematic for us. Just pass None or whatever you want

ricardoV94 · 2024-04-15T09:20:12Z

tests/graph/test_op.py

+    res_single = single_op(x, name=op_name)
+    assert res_single.name == op_name


This may be more readable. Also the name we are giving is not so much an op_name but a var_name.

Suggested change

res_single = single_op(x, name=op_name)

assert res_single.name == op_name

res_single = single_op(x, name="test_name")

assert res_single.name == "test_name"

Also test that by default name is None?

ricardoV94 · 2024-04-15T09:21:15Z

tests/graph/test_op.py

+        def make_node(self, *inputs):
+            outputs = [dummy_variable("a"), dummy_variable("b")]
+            return Apply(self, list(inputs), outputs)


Suggested change

def make_node(self, *inputs):

outputs = [dummy_variable("a"), dummy_variable("b")]

return Apply(self, list(inputs), outputs)

def make_node(self, input):

inputs = [input]

outputs = [input.type(), input.type()]

return Apply(self, inputs, outputs)

ricardoV94 · 2024-04-15T09:21:40Z

tests/graph/test_op.py

+    res = multi_op(x, name=op_name)
+    for i, r in enumerate(res):
+        assert r.name == f"{op_name}_{i}"


Suggested change

res = multi_op(x, name=op_name)

for i, r in enumerate(res):

assert r.name == f"{op_name}_{i}"

res = multi_op(x, name="test_name")

for i, r in enumerate(res):

assert r.name == f"test_name_{i}"

ricardoV94 · 2024-04-15T09:22:53Z

tests/graph/test_op.py

+    multi_op = MultiOutOp()
+    multi_op.default_output = 1


A bit of a nitpick, but more "realistic" if you make the default_output something defined when a class is initialized. Just need to define a __init__ method that assigns it to self after calling super().init.

Suggested change

multi_op = MultiOutOp()

multi_op.default_output = 1

multi_op = MultiOutOp(default_output=1)

You can also parametrize the number of outputs, with a multi_output=False|True. That way you can use a single Test Op class, which reduces a bit the number lines of code in this test. After storing it in init, in made_node you can do:

if self.multi_output: outputs = [input.type(), input.type()] else: outputs = [input.type()]

ricardoV94 · 2024-04-15T09:24:05Z

tests/graph/test_op.py

+        def make_node(self, *inputs):
+            outputs = [dummy_variable("a")]
+            return Apply(self, list(inputs), outputs)


Traditionally, an Op would not assign a default name in make_node

Suggested change

def make_node(self, *inputs):

outputs = [dummy_variable("a")]

return Apply(self, list(inputs), outputs)

def make_node(self, input):

inputs = [input]

outputs = [input.type()]

return Apply(self, inputs, outputs)

ricardoV94 · 2024-04-15T09:25:59Z

Functionality looks great, just some test suggestions

ricardoV94 · 2024-04-15T09:26:56Z

tests/graph/test_op.py

@@ -232,3 +232,55 @@ def perform(self, *_):

    x = pt.TensorType(dtype="float64", shape=(1,))("x")
    assert SomeOp()(x).type == pt.dvector
+
+
+def test_op_name():


Suggested change

def test_op_name():

def test_call_name():

ricardoV94 · 2024-04-15T09:29:06Z

tests/graph/test_op.py

+    multi_op = MultiOutOp()
+    multi_op.default_output = 1
+    res = multi_op(x, name=op_name)
+    assert res.name == op_name


Test that res.owner.outputs that are not the default_output still have name is None

codecov · 2024-04-19T20:30:32Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 80.76%. Comparing base (f97d9ea) to head (1c99e01).
Report is 18 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #693      +/-   ##
==========================================
- Coverage   80.83%   80.76%   -0.08%     
==========================================
  Files         162      162              
  Lines       46830    46713     -117     
  Branches    11447    11426      -21     
==========================================
- Hits        37857    37729     -128     
- Misses       6710     6735      +25     
+ Partials     2263     2249      -14

Files	Coverage Δ
pytensor/graph/op.py	`87.89% <100.00%> (+0.39%)`	⬆️

... and 22 files with indirect coverage changes

ricardoV94 · 2024-04-20T06:17:58Z

Thanks @HarshvirSandhu !

HarshvirSandhu added 2 commits April 8, 2024 00:40

Add name kwarg to Op.__call__

bb12ed0

Fix Index out of bounds error

568334b

ricardoV94 reviewed Apr 8, 2024

View reviewed changes

pytensor/graph/op.py Outdated Show resolved Hide resolved

HarshvirSandhu added 3 commits April 11, 2024 23:38

Fix mypy error

3e6379e

Fix failing tests

182b7ae

Fix ruff format

64c121a

Dhruvanshu-Joshi reviewed Apr 12, 2024

View reviewed changes

pytensor/graph/op.py Outdated Show resolved Hide resolved

Remove print statement

1755294

twiecki approved these changes Apr 12, 2024

View reviewed changes

ricardoV94 requested changes Apr 12, 2024

View reviewed changes

Dhruvanshu-Joshi mentioned this pull request Apr 12, 2024

Add pre-commit hook to avoid print statements. #709

Closed

Modify names for multi-output nodes

31f51eb

ricardoV94 requested changes Apr 13, 2024

View reviewed changes

Ignore naming Ops without name kwarg

1470b1b

ricardoV94 reviewed Apr 13, 2024

View reviewed changes

Modify test function

183dd3f

HarshvirSandhu requested a review from ricardoV94 April 14, 2024 19:51

ricardoV94 reviewed Apr 15, 2024

View reviewed changes

Parametrize test and use single dummy op

1c99e01

ricardoV94 approved these changes Apr 20, 2024

View reviewed changes

ricardoV94 merged commit 86bc1d2 into pymc-devs:main Apr 20, 2024

		node = self.make_node(inputs, *kwargs)
		if isinstance(node.outputs, list):

		res_single = single_op(x, name=op_name)
		assert res_single.name == op_name

-        def make_node(self, *inputs):
-            outputs = [dummy_variable("a"), dummy_variable("b")]
-            return Apply(self, list(inputs), outputs)
+        def make_node(self, input):
+            inputs = [input]
+            outputs = [input.type(), input.type()]
+            return Apply(self, inputs, outputs)

	multi_op = MultiOutOp()
	multi_op.default_output = 1
	multi_op = MultiOutOp(default_output=1)

Add name kwarg to Op.__call__ #693

Add name kwarg to Op.__call__ #693

Uh oh!

Conversation

HarshvirSandhu commented Apr 7, 2024

Description

Related Issue

Checklist

Type of change

Uh oh!

Uh oh!

Uh oh!

ricardoV94 Apr 12, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Apr 13, 2024

Uh oh!

HarshvirSandhu commented Apr 13, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

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!

ricardoV94 Apr 15, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Apr 15, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Apr 19, 2024

Codecov Report

Uh oh!

ricardoV94 commented Apr 20, 2024

Uh oh!

Uh oh!

Add name kwarg to Op.call #693

Add name kwarg to Op.call #693

ricardoV94 Apr 12, 2024 •

edited

Loading

ricardoV94 Apr 15, 2024 •

edited

Loading