finch-tensor / finch-tensor-python

Sparse and Structured Tensor Programming in Python
MIT License
8 stars 3 forks source link

API: Add `eye` function #32

Closed mtsokol closed 6 months ago

mtsokol commented 6 months ago

Hi @willow-ahrens @hameerabbasi,

This PR adds eyeArray API function.

@willow-ahrens I think that as a follow-up I could add format argument but, for COO:

using Finch
tns = Tensor(SparseCOO{2}(Element(0)), 3, 4)
@finch begin
    for i=_, j=_
        if i == j
           tns[i, j] += 1.0
        end
    end
end

it fails with:

ERROR: Finch.FinchProtocolError("cannot modify Finch.VirtualSparseCOOLevel in place (forgot to declare with .= ?)")
willow-ahrens commented 6 months ago

Declare it first with tns .= 0

willow-ahrens commented 6 months ago

Do you think we should just add speye and spdiags to finch? I'll work on that this evening.

mtsokol commented 6 months ago

eye is pretty simple to implement in finch-tensor, but sure - Finch functions that follow:

mtsokol commented 6 months ago

We could have format=... argument there? Or always return coo

willow-ahrens commented 6 months ago

I just realized that speye and spdiags are classic matlab functions, so we should probably also do it in Finch.jl

mtsokol commented 6 months ago

Declare it first with tns .= 0

It helped - thank you! For now we have format=... keyword argument that supports dense and coo by building an expression that is evaluated.

willow-ahrens commented 6 months ago

Sure! Maybe we should merge this and create an issue in the Finch repo for spdiags and speye

mtsokol commented 6 months ago

Sure!

mtsokol commented 6 months ago

I created an issue. The PR is ready I think.