I believe the following is an implementation using OpFromGraph:

n = at.iscalar('n')
m = at.iscalar('m')
k = at.iscalar('k')

i = at.switch(k >= 0, k, -k * m)
eye = at.zeros(n * m)
eye = at.set_subtensor(eye[i::m+1], 1).reshape((n, m))
eye = at.set_subtensor(eye[m-k:, :], 0)

Eye = aesara.compile.builders.OpFromGraph([n, m, k], [eye])

The existing eye function in aesara.tensor could then be modified as follows:

def eye(n, m=None, k=0, dtype=None):
    """Return a 2-D array with ones on the diagonal and zeros elsewhere.
    Parameters
    ----------
    n : int
        Number of rows in the output.
    m : int, optional
        Number of columns in the output. If None, defaults to `N`.
    k : int, optional
        Index of the diagonal: 0 (the default) refers to the main diagonal,
        a positive value refers to an upper diagonal, and a negative value
        to a lower diagonal.
    dtype : data-type, optional
        Data-type of the returned array.
    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.
    """
    if dtype is None:
        dtype = aesara.config.floatX
    if m is None:
        m = n
    return Eye(n, m, k).astype(dtype)

I can add some tests to check corner cases and submit this as a pull request if it looks like I'm barking up the right tree? Where would I add the code to create the OpFromGraph, just floating inside aesara.tensor above def eye, or is there a more organized place to put it?

Uh oh!

Convert Eye to a COp or implement it in terms of existing COps #1177

Description

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Convert `Eye` to a `COp` or implement it in terms of existing `COp`s #1177