Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Implement at.eye using existing Ops #1217

Open
wants to merge 26 commits into
base: main
Choose a base branch
from
Open

Implement at.eye using existing Ops #1217

Show file tree
Hide file tree
Changes from 21 commits
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
6617b69
added Ops for solve_discrete_lyapunov and solve_continuous_lyapunov
jessegrabowski Jun 29, 2022
34acf8f
ran pre-commit
jessegrabowski Jun 29, 2022
f28cee4
add `method` to `SolveDiscreteLyapunov` `__props__`
jessegrabowski Jun 30, 2022
10472ba
Remove shape checks from `SolveDiscreteLyapunov` and `SolveContinuous…
jessegrabowski Jun 30, 2022
0ec6c10
Merge branch 'solve_lyapunov' of github.com:jessegrabowski/aesara int…
jessegrabowski Jun 30, 2022
bf445ab
Update signature of `solve_discrete_lyapunov` and `solve_continuous_l…
jessegrabowski Jun 30, 2022
8a9da62
Update signature of `solve_discrete_lyapunov` and `solve_continuous_l…
jessegrabowski Jun 30, 2022
1738058
Delete a scratchpad file
jessegrabowski Jun 30, 2022
5c2442b
Add a direct aesara solution via `at.solve` for `solve_discrete_lyapu…
jessegrabowski Sep 27, 2022
e0b5d2e
add `method` to `SolveDiscreteLyapunov` `__props__`
jessegrabowski Jun 30, 2022
a14303a
Update signature of `solve_discrete_lyapunov` and `solve_continuous_l…
jessegrabowski Jun 30, 2022
4c095b6
Update signature of `solve_discrete_lyapunov` and `solve_continuous_l…
jessegrabowski Jun 30, 2022
810f5d9
Delete a scratchpad file
jessegrabowski Jun 30, 2022
837ec91
Add a direct aesara solution via `at.solve` for `solve_discrete_lyapu…
jessegrabowski Sep 27, 2022
2487597
Merge branch 'solve_lyapunov' of github.com:jessegrabowski/aesara int…
jessegrabowski Sep 27, 2022
2861abd
Rewrite function `eye` using `at.zeros` and `at.set_subtensor`
jessegrabowski Sep 27, 2022
4ac63b3
remove changes unrelated to eye
jessegrabowski Sep 27, 2022
af2dc7d
remove changes unrelated to eye
jessegrabowski Sep 27, 2022
94ce70d
remove changes unrelated to eye
jessegrabowski Sep 27, 2022
29991e3
remove changes unrelated to eye
jessegrabowski Sep 27, 2022
b4d9b7e
Merge branch 'main' into aesara_native_eye
jessegrabowski Sep 27, 2022
4c2a622
reverted unnecessary changes to `test_eye`
jessegrabowski Sep 27, 2022
b823eec
Merge branch 'aesara_native_eye' of github.com:jessegrabowski/aesara …
jessegrabowski Sep 27, 2022
680179e
Merge branch 'aesara-devs:main' into aesara_native_eye
jessegrabowski Oct 1, 2022
5ece747
Add typing information to arguments in function signature
jessegrabowski Oct 1, 2022
741a73c
Merge branch 'aesara_native_eye' of github.com:jessegrabowski/aesara …
jessegrabowski Oct 1, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 17 additions & 8 deletions aesara/tensor/basic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1302,17 +1302,26 @@ def eye(n, m=None, k=0, dtype=None):

Returns
-------
ndarray of shape (N,M)
An array where all elements are equal to zero, except for the `k`-th
diagonal, whose values are equal to one.

aesara tensor of shape (N,M)
A symbolic tensor representing a matrix where all elements are equal to zero,
except for the `k`-th diagonal, whose values are equal to one.
brandonwillard marked this conversation as resolved.
Show resolved Hide resolved
"""
if dtype is None:
dtype = config.floatX

if m is None:
m = n
localop = Eye(dtype)
return localop(n, m, k)
if dtype is None:
dtype = aesara.config.floatX

n = as_tensor_variable(n)
m = as_tensor_variable(m)
k = as_tensor_variable(k)

i = switch(k >= 0, k, -k * m)
brandonwillard marked this conversation as resolved.
Show resolved Hide resolved
eye = zeros(n * m, dtype=dtype)
eye = aesara.tensor.subtensor.set_subtensor(eye[i :: m + 1], 1).reshape((n, m))
eye = aesara.tensor.subtensor.set_subtensor(eye[m - k :, :], 0)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need the last set_subtensor (i.e. that appears to fill with 0s if we've already used zeros? I think I might know why, but it would be nice to have only one set_subtensor.

Copy link
Contributor Author

@jessegrabowski jessegrabowski Sep 30, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The second set_subtenor is to prevent the diagonal from "wrapping around" the matrix in cases when m > n or when k != 0. This one-liner almost works:

at.subtensor.set_subtensor(eye[i:(m - k)*m:m + 1], 1).reshape((n, m))

But it fails when k < -n or k > m. In these cases I guess an error could be raised (why is the user asking for an offset greater than the width of the matrix?), but numpy returns the zero matrix for this case and I wanted to be consistent.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One way to match numpy and avoid the second set_subtensor would be to check the conditions where the zero matrix should be returned and use multiplication:

    eye = at.subtensor.set_subtensor(eye[i:(m - k)*m:m + 1], 1).reshape((n, m))
    eye *= (1 - at.gt(k, m)) * (1 - at.lt(k, -n))

But this was about 5µs slower than the two set_subtensor method when I ran %timeit. Perhaps it's too cute?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here's an implementation that uses scalars and Composite:

import numpy as np
import aesara
import aesara.tensor as at
from aesara.scalar import int64, switch, Composite


n = int64("n")
m = int64("m")
k = int64("k")

i = switch(k >= 0, k, -k * m)

eye_sc = at.zeros(n * m, dtype="float64")
mkm = (m - k) * m
eye_sc = at.subtensor.set_subtensor(eye_sc[i : mkm : m + 1], 1).reshape((n, m))

eye_sc_fn = aesara.function([n, m, k], [eye_sc])

aesara.dprint(eye_sc_fn)
# Reshape{2} [id A] 16
#  |IncSubtensor{InplaceSet;int64:int64:int64} [id B] 15
#  | |Reshape{2} [id C] 14
#  | | |IncSubtensor{InplaceSet;int64:int64:int64} [id D] 13
#  | | | |Alloc [id E] 11
#  | | | | |TensorConstant{0.0} [id F]
#  | | | | |TensorFromScalar [id G] 8
#  | | | |   |mul [id H] 2
#  | | | |     |n [id I]
#  | | | |     |m [id J]
#  | | | |TensorConstant{1} [id K]
#  | | | |Switch [id L] 12
#  | | | | |GE [id M] 6
#  | | | | | |k [id N]
#  | | | | | |ScalarConstant{0} [id O]
#  | | | | |k [id N]
#  | | | | |mul [id P] 10
#  | | | |   |neg [id Q] 5
#  | | | |   | |k [id N]
#  | | | |   |m [id J]
#  | | | |mul [id R] 9
#  | | | | |sub [id S] 4
#  | | | | | |m [id J]
#  | | | | | |k [id N]
#  | | | | |m [id J]
#  | | | |add [id T] 3
#  | | |   |m [id J]
#  | | |   |ScalarConstant{1} [id U]
#  | | |MakeVector{dtype='int64'} [id V] 7
#  | |   |TensorFromScalar [id W] 1
#  | |   | |n [id I]
#  | |   |TensorFromScalar [id X] 0
#  | |     |m [id J]
#  | |TensorConstant{1} [id K]
#  | |Switch [id L] 12
#  | |mul [id R] 9
#  | |add [id T] 3
#  |MakeVector{dtype='int64'} [id V] 7

%timeit eye_sc_fn(10, 10, 1)
# 66.5 µs ± 1.42 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

# Manually construct composite `Op`s for the scalar operations
i_comp_op = Composite([n, m, k], [i])
i_comp = i_comp_op(n, m, k)

mkm_comp_op = Composite([m, k], [mkm])
mkm_comp = mkm_comp_op(m, k)

eye_sc_comp = at.zeros(n * m, dtype="float64")
eye_sc_comp = at.subtensor.set_subtensor(eye_sc_comp[i_comp : mkm_comp : m + 1], 1).reshape((n, m))

eye_sc_comp_fn = aesara.function([n, m, k], [eye_sc_comp])

aesara.dprint(eye_sc_comp_fn)
# Reshape{2} [id A] 10
#  |IncSubtensor{InplaceSet;int64:int64:int64} [id B] 9
#  | |Alloc [id C] 8
#  | | |TensorConstant{0.0} [id D]
#  | | |TensorFromScalar [id E] 7
#  | |   |mul [id F] 6
#  | |     |n [id G]
#  | |     |m [id H]
#  | |TensorConstant{1} [id I]
#  | |Composite{Switch(GE(i2, 0), i2, ((-i2) * i1))} [id J] 5
#  | | |n [id G]
#  | | |m [id H]
#  | | |k [id K]
#  | |Composite{((i0 - i1) * i0)} [id L] 4
#  | | |m [id H]
#  | | |k [id K]
#  | |add [id M] 3
#  |   |m [id H]
#  |   |ScalarConstant{1} [id N]
#  |MakeVector{dtype='int64'} [id O] 2
#    |TensorFromScalar [id P] 1
#    | |n [id G]
#    |TensorFromScalar [id Q] 0
#      |m [id H]

%timeit eye_sc_comp_fn(10, 10, 1)
# 60.9 µs ± 2.53 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

Copy link
Member

@brandonwillard brandonwillard Sep 30, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Before I forget, using inc_subtensor with set_instead_of_inc=True and ignore_duplicates=True might be faster. Only applies when AdvancedIncSubtensor is used, and it isn't in these cases.

Copy link
Member

@brandonwillard brandonwillard Sep 30, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We also need an issue for a rewrite that replaces scalar-only sub-graphs with Composites.

Done: #1225

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, Switch is quite literally the scalar version of IfElse—at least in terms of its Op.c_code implementation. If anything, we should make invocations of ifelse return Switches when the arguments are ScalarTypeed.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here's the issue for ifelse: #1226.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, there's no scalar Op for max and min?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Disregard, obviously this can be done using scalar switch.


return eye


def identity_like(x, dtype: Optional[Union[str, np.generic, np.dtype]] = None):
Expand Down
7 changes: 2 additions & 5 deletions tests/tensor/test_basic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -822,11 +822,8 @@ def check(dtype, N, M_=None, k=0):
# allowed.
if M is None and config.mode in ["DebugMode", "DEBUG_MODE"]:
M = N
N_symb = iscalar()
M_symb = iscalar()
k_symb = iscalar()
f = function([N_symb, M_symb, k_symb], eye(N_symb, M_symb, k_symb, dtype=dtype))
result = f(N, M, k)

result = eye(N, M, k, dtype).eval()
jessegrabowski marked this conversation as resolved.
Show resolved Hide resolved
assert np.allclose(result, np.eye(N, M_, k, dtype=dtype))
assert result.dtype == np.dtype(dtype)

Expand Down