Fix typing in subtensor module #823

michaelosthege · 2024-06-16T16:15:40Z

Description

I started fixing type hints here and there, and then focused on the pytensor.tensor.subtensor module.

It was tricky, but by adding some overloads (to help mypy) and refactoring one function, I managed to clear errors in that file.

The only remaining error is the Cannot determine type thing for an imported dispatched function.

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

codecov · 2024-06-16T16:43:11Z

Codecov Report

Attention: Patch coverage is 82.00000% with 9 lines in your changes missing coverage. Please review.

Project coverage is 80.98%. Comparing base (7a00b88) to head (acd537b).
Report is 152 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/subtensor.py	79.54%	1 Missing and 8 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #823      +/-   ##
==========================================
- Coverage   80.99%   80.98%   -0.02%     
==========================================
  Files         169      169              
  Lines       46939    46961      +22     
  Branches    11467    11478      +11     
==========================================
+ Hits        38019    38032      +13     
- Misses       6713     6714       +1     
- Partials     2207     2215       +8

Files with missing lines	Coverage Δ
pytensor/graph/basic.py	`88.61% <ø> (ø)`
pytensor/tensor/random/op.py	`93.78% <100.00%> (ø)`
pytensor/tensor/random/utils.py	`100.00% <100.00%> (ø)`
pytensor/tensor/shape.py	`92.68% <100.00%> (ø)`
pytensor/tensor/type.py	`94.54% <100.00%> (ø)`
pytensor/tensor/subtensor.py	`89.27% <79.54%> (-0.62%)`	⬇️

pytensor/tensor/subtensor.py

michaelosthege · 2024-06-17T05:45:49Z

the coverage ❌ is due to the overload signatures - can't do anything about it

ricardoV94

Had a look, have some doubts about a few of the changes, but the directions seems great.

In addition, if you want to remove the lines about sys.maxsize from the giant canonical slice function feel free to do it, as we are not using that for sure (and there's even an inline comment in one of its uses, saying so)

ricardoV94 · 2024-06-17T08:26:07Z

pytensor/tensor/subtensor.py

+                else:
+                    sslice = vlit
+            except NotScalarConstantError:
+                raise ValueError(f"Slice {theslice} is not a supported slice type.")


~~Couldn't it be a slice with symbolic start/step/stop at this point?~~

From reading the code it seems you changed the meaning of the function to fail with certain symbolic inputs, such as slice(pt.lscalar("start"), pt.lscalar("stop"), pt.lscalar("step")) or a scalar indexing pt.lscalar("idx"), where it was accommodating them before

Okay this function does not expect symbolic slices (or at least it was never testing them), so the new changes shoudn't either?

If it's not a slice it must be a scalar index, and you shouldn't have to worry about the slice case showing up again here, the old code was certainly not worried about that case. If it should handle symbolic slices (i.e, SliceVariable) types, then we also need tests for those.

It wasn't testing them, but as_index_literal already took them apart into a native slice:

pytensor/pytensor/tensor/subtensor.py

Line 172 in efa845a

idx = slice(*idx.owner.inputs)

Consequently, get_canonical_form_slice handled them already..

That still means this function only needs to handle python slices, and the inner code can be simplified

ricardoV94 · 2024-06-17T08:32:00Z

pytensor/tensor/subtensor.py

+                if vlit is None:
+                    return slice(0, length, 1), 1


Are you sure? Isn't None a newaxis, which shouldn't correspond to any slice?

Also unclear whether this function is expected to handle None. Like SliceVariable, those only show up in the index for AdvancedSubtensor Ops. For basic Subtensor Ops, newaxis are eagerly converted to expand_dims, and never show as inputs of the Op

From the docstring I would say the function is not supposed to do None → None things.
There were no tests for that either, and we would create a bigger mess by hinting a possibly None return type.

Actually, I think we should raise ValueError in this situation.

ricardoV94 · 2024-06-17T08:39:45Z

Regarding the graphs being complicated, there's a related issue #112

pytensor/tensor/subtensor.py

ricardoV94 · 2024-06-17T08:48:38Z

pytensor/tensor/subtensor.py

-            value = theslice
-
-        value = switch(lt(value, 0), (value + length), value)
+        # Try to extract a scalar slice and return it already.


~~This branch may be more clean if it clearly tries to not handle slices, by checking not isinstance(theslice, slice | SliceVariable)?~~ I don't know if newaxis are supposed to make it to this function, I had a question below about it.

And I would call these scalar indexing instead of slices, not scalar slices?

~~Then after this branch you can force the SliceVariable to be a normal slice by doing as_index_literal or slice(*sslice.owner.inputs)~~

Edit: After reading the pre-existig code and tests I am not sure this function is supposed to handle SliceVariables, so ignore that part of the remark

It can do that, but I don't think we should "advertise" it in the type hints. It's covered by the Variable hint already.

pytensor/tensor/subtensor.py

tests/tensor/test_subtensor.py

ricardoV94 · 2024-06-17T08:55:47Z

tests/tensor/test_subtensor.py



 class TestGetCanonicalFormSlice:
+    def test_none_constant(self):


Very unsure about this. None != slice(None) in indexing

Yes, I changed this into raising a ValueError, because addition of a new axis something other than canonical slicing.

See comment above too.

ricardoV94 · 2024-06-17T08:56:26Z

tests/tensor/test_subtensor.py

+    def test_symbolic_slice(self):
+        idx = make_slice(slice(3, 7, 2))
+        assert not isinstance(idx, slice)
+        res = get_canonical_form_slice(idx, 10)


Nitpick suggestion for these kinds of tests

Suggested change

res = get_canonical_form_slice(idx, 10)

canonical_slice, direction = get_canonical_form_slice(idx, 10)

Too scared of black's line wrapping wrath

pytensor/tensor/subtensor.py

ricardoV94 · 2024-06-17T09:13:00Z

tests/tensor/test_subtensor.py

+        assert res[1] == 1
+
+    def test_symbolic_slice(self):
+        idx = make_slice(slice(3, 7, 2))


This is the part I was confused. I thought the function was meant to handle these types originally, but after more careful reading, don't think it was?

If it was (or became) meant to handle them, then it should also handle make_slice, with symbolic inputs.

No it can handle variable and constant symbolic slices too 😎

pytensor/tensor/subtensor.py

ricardoV94 · 2024-06-17T09:23:21Z

pytensor/tensor/subtensor.py

+
+
+@overload
+def as_index_literal(idx: Constant) -> int | np.integer | None: ...


You can be more specific by using NoneConst, SliceConst, Tensor|ScalarConstant and SliceVariable, Tensor|ScalarVariable in the overloaded signatures

TensorContants are already a subclass of TensorVariable (so you don't need that distinction at the signature level), but the other types ScalarConstant | SliceConstant may not be. You can make them if you want.

This is the PR where we made some of those subclasses: #628

Thanks for point out that these types exist. With them, I was able to simplify the function, and among other things avoid the getattr too.

I will not reply to your other comments on as_index_literal, since they no longer apply.

pytensor/tensor/subtensor.py

michaelosthege

Thanks @ricardoV94, your suggestions helped to further simplify the code.

~~On my machine, these changes accelerated the test_subtensor.py suite by almost 10 seconds, even though new test cases were added(!)~~

pytensor/tensor/subtensor.py

michaelosthege · 2024-06-19T19:48:36Z

pytensor/tensor/subtensor.py

+
+
+@overload
+def as_index_literal(idx: Constant) -> int | np.integer | None: ...


Thanks for point out that these types exist. With them, I was able to simplify the function, and among other things avoid the getattr too.

I will not reply to your other comments on as_index_literal, since they no longer apply.

michaelosthege · 2024-06-19T20:19:40Z

pytensor/tensor/subtensor.py

-            value = theslice
-
-        value = switch(lt(value, 0), (value + length), value)
+        # Try to extract a scalar slice and return it already.


It can do that, but I don't think we should "advertise" it in the type hints. It's covered by the Variable hint already.

michaelosthege · 2024-06-19T20:27:26Z

pytensor/tensor/subtensor.py

+                if vlit is None:
+                    return slice(0, length, 1), 1


From the docstring I would say the function is not supposed to do None → None things.
There were no tests for that either, and we would create a bigger mess by hinting a possibly None return type.

Actually, I think we should raise ValueError in this situation.

michaelosthege · 2024-06-19T20:35:15Z

pytensor/tensor/subtensor.py

+                else:
+                    sslice = vlit
+            except NotScalarConstantError:
+                raise ValueError(f"Slice {theslice} is not a supported slice type.")


It wasn't testing them, but as_index_literal already took them apart into a native slice:

pytensor/pytensor/tensor/subtensor.py

Line 172 in efa845a

idx = slice(*idx.owner.inputs)

Consequently, get_canonical_form_slice handled them already..

michaelosthege · 2024-06-19T20:37:51Z

tests/tensor/test_subtensor.py



 class TestGetCanonicalFormSlice:
+    def test_none_constant(self):


Yes, I changed this into raising a ValueError, because addition of a new axis something other than canonical slicing.

See comment above too.

tests/tensor/test_subtensor.py

michaelosthege · 2024-06-19T21:11:48Z

tests/tensor/test_subtensor.py

+        assert res[1] == 1
+
+    def test_symbolic_slice(self):
+        idx = make_slice(slice(3, 7, 2))


No it can handle variable and constant symbolic slices too 😎

michaelosthege · 2024-06-19T21:12:51Z

tests/tensor/test_subtensor.py

+    def test_symbolic_slice(self):
+        idx = make_slice(slice(3, 7, 2))
+        assert not isinstance(idx, slice)
+        res = get_canonical_form_slice(idx, 10)


Too scared of black's line wrapping wrath

ricardoV94

Looking cleaner, but I still think we can simplify, because we are doing more now than the function was doing before

ricardoV94 · 2024-06-21T08:46:35Z

pytensor/tensor/subtensor.py

        )

+    if not isinstance(idx, Variable):
+        raise NotScalarConstantError()


Suggested change

raise NotScalarConstantError()

raise TypeError()

pytensor/tensor/subtensor.py

ricardoV94 · 2024-06-21T08:50:43Z

pytensor/tensor/subtensor.py


+@overload
+def get_canonical_form_slice(
+    theslice: slice | SliceConstant,


It's not just Constants right?

Suggested change

theslice: slice | SliceConstant,

theslice: slice | SliceVariable,

Yes, but there's no SliceVariable type at the moment

Oh strange. So there are Variables with type SliceType but no specific class like TensorVariable or ScalarVariable?

pytensor/tensor/subtensor.py

ricardoV94 · 2024-06-21T08:52:23Z

pytensor/tensor/subtensor.py

 def get_canonical_form_slice(
-    theslice: slice | Variable, length: Variable
-) -> tuple[Variable, int]:
+    theslice: slice | Variable,


Use the union of the specific types we overloaded above? Variable allows more than SliceVariable | ScalarVariable | TensorVariable but our function does not

The Variable input covers the slice-variable for which we don't have a dedicated type.

We should create one then

ricardoV94 · 2024-06-21T08:53:40Z

pytensor/tensor/subtensor.py

-
-        value = switch(lt(value, 0), (value + length), value)
+    # Convert the two symbolic slice types into a native slice
+    if isinstance(theslice, SliceConstant):


I still think this function is not supposed to handle non-python slices (be it constant or not).

That allows us to simplify our logic and type hints quite a lot?

As long as symbolic slices exist, I think this function should support them. How else should they be handled?

For now, I'd just like to get the typing fixed.

This function does not seem to be intended for SliceTypes. I wrote elsewhere but there are two ways of representing slices in PyTensor: basic index operations don't have explicit slice inputs, they just take as inputs the start/stop/step variables that define a slice.

Advanced indexing operations instead have explicit slice variables of SliceVariable type as inputs. This function was seemingly not meant to work with those types and unless that's why you went into this refactor there's no reason to extend its behavior to handle the slice types that are ony used by advanced indexing operations.

Does that make sense?

ricardoV94 · 2024-06-21T08:55:20Z

pytensor/tensor/subtensor.py

+            sslice = theslice
+        elif isinstance(theslice, TensorVariable) and theslice.ndim == 0:
+            sslice = theslice
+        else:


Just raise a TypeError, no other types should show up

There are almost a hundred places where NotScalarConstantError is caught. Replacing it with a generic TypeError appears very risky..

It's being misused here, why do you want to raise it?

ricardoV94 · 2024-06-21T08:55:47Z

pytensor/tensor/subtensor.py

+                else:
+                    sslice = vlit
+            except NotScalarConstantError:
+                raise ValueError(f"Slice {theslice} is not a supported slice type.")


That still means this function only needs to handle python slices, and the inner code can be simplified

ricardoV94 · 2024-06-21T08:57:19Z

tests/tensor/test_subtensor.py

+
+    def test_constant_slice(self):
+        idx = as_symbolic_slice(slice(3, 7, 2))
+        assert isinstance(idx, SliceConstant)


I don't think we need to support non python slices. The code before wasn't, so there is no specified need?

Why are there symbolic slices to begin with?

If we don't need them, we should remove them. If we need them, I think the functions dealing with slices should deal with them too.

Replied here: #823 (comment)

They are used in a different class of Ops

pytensor/tensor/subtensor.py

ricardoV94 · 2024-06-25T05:59:14Z

Thanks @michaelosthege

michaelosthege self-assigned this Jun 16, 2024

michaelosthege force-pushed the subtensor-typing branch from a9798f2 to 2b77374 Compare June 16, 2024 16:20

michaelosthege force-pushed the subtensor-typing branch from 2b77374 to 6200b05 Compare June 16, 2024 17:40

michaelosthege commented Jun 16, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

michaelosthege requested review from Armavica and ricardoV94 June 16, 2024 18:04

michaelosthege added maintenance refactor labels Jun 16, 2024

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

tests/tensor/test_subtensor.py Show resolved Hide resolved

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

pytensor/tensor/subtensor.py Show resolved Hide resolved

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jun 17, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

michaelosthege commented Jun 19, 2024

View reviewed changes

michaelosthege changed the title ~~Fiy typing in subtensor module~~ Fix typing in subtensor module Jun 19, 2024

ricardoV94 requested changes Jun 21, 2024

View reviewed changes

ricardoV94 reviewed Jun 23, 2024

View reviewed changes

pytensor/tensor/subtensor.py Show resolved Hide resolved

ricardoV94 reviewed Jun 24, 2024

View reviewed changes

pytensor/tensor/subtensor.py Outdated Show resolved Hide resolved

michaelosthege added 2 commits June 24, 2024 19:59

Fix type hints here and there

2ed85dc

Refactor get_canonical_form_slice to fix subtensor typing

acd537b

michaelosthege force-pushed the subtensor-typing branch from 67c8df9 to acd537b Compare June 24, 2024 22:11

ricardoV94 approved these changes Jun 25, 2024

View reviewed changes

ricardoV94 merged commit f9dfe70 into pymc-devs:main Jun 25, 2024

michaelosthege deleted the subtensor-typing branch June 25, 2024 06:35



		class TestGetCanonicalFormSlice:
		def test_none_constant(self):

	res = get_canonical_form_slice(idx, 10)
	canonical_slice, direction = get_canonical_form_slice(idx, 10)



		@overload
		def as_index_literal(idx: Constant) -> int \| np.integer \| None: ...

	theslice: slice \| SliceConstant,
	theslice: slice \| SliceVariable,

Fix typing in subtensor module #823

Fix typing in subtensor module #823

Uh oh!

Conversation

michaelosthege commented Jun 16, 2024

Description

Checklist

Type of change

Uh oh!

codecov bot commented Jun 16, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

michaelosthege commented Jun 17, 2024

Uh oh!

ricardoV94 left a comment • 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!

ricardoV94 Jun 17, 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!

ricardoV94 Jun 17, 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!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jun 17, 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 Jun 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ricardoV94 Jun 17, 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!

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!

Uh oh!

Uh oh!

ricardoV94 Jun 17, 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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ricardoV94 Jun 17, 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

codecov bot commented Jun 16, 2024 •

edited

Loading

ricardoV94 left a comment •

edited

Loading

ricardoV94 Jun 17, 2024 •

edited

Loading

ricardoV94 Jun 17, 2024 •

edited

Loading

ricardoV94 Jun 17, 2024 •

edited

Loading

ricardoV94 commented Jun 17, 2024 •

edited

Loading

ricardoV94 Jun 17, 2024 •

edited

Loading

ricardoV94 Jun 17, 2024 •

edited

Loading

ricardoV94 Jun 17, 2024 •

edited

Loading

michaelosthege left a comment •

edited

Loading