Tracking issue: Add Array API standard support

[mtsokol]

Hi @willow-ahrens @hameerabbasi,

This issue is meant to track progress of implementing Array API standard for finch-tensor.

I thought that we could try adding short notes to bullet-points, saying which Finch.jl functions should be called to implement given entry. I think we already had some ideas during one of our first calls.

Array API: https://data-apis.org/array-api/latest/index.html

Backlog

main namespace

• [x] astype - https://github.com/willow-ahrens/finch-tensor/pull/15
• [x] element-wise ops (add, multiply, cos, ...) - https://github.com/willow-ahrens/finch-tensor/pull/17 (partially...)
• [x] reductions (xp.prod, xp.sum) - jl.sum and jl.prod, also just jl.reduce - https://github.com/willow-ahrens/finch-tensor/pull/17
• [x] matmul - implemented with finch.tensordot for non-stacked input. Should be rewritten with jl.mul / Finch einsum.
• [x] tensordot - finch.tensordot - https://github.com/willow-ahrens/finch-tensor/pull/22
• [x] where - jl.broadcast(jl.ifelse, cond, a, b) - https://github.com/willow-ahrens/finch-tensor/pull/30
• [ ] argmin/argmax - jl.argmin (bug willow if this isn't implemented already)
• [ ] take - jl.getindex
• [x] nonzero - this is an eager function, but it is implemented as ffindnz(arr) - https://github.com/willow-ahrens/finch-tensor/pull/30
• [x] creation functions: asarray, ones, full, full_like, ... - finch.Tensor constructor, as well as jl.copyto!(arr, jl.broadcasted(Scalar(1)), as well as changing the default of the tensor with Tensor(Dense(Element(1.0))). We may need to distinguish some of these. https://github.com/willow-ahrens/finch-tensor/pull/28, https://github.com/willow-ahrens/finch-tensor/pull/32
• [ ] stats functions: max, mean, min, std, var
• [ ] set functions: unique_all, unique_counts, unique_inverse, unique_values
• [x] all, any
• [ ] concat
• [ ] expand_dims
• [ ] flip
• [ ] reshape
• [ ] roll
• [ ] squeeze
• [ ] stack
• [ ] argsort/sort
• [ ] broadcast_arrays
• [ ] broadcast_to
• [x] can_cast/finfo/iinfo/result_type
• [x] bitwise ops: bitwise_and/bitwise_left_shift/bitwise_invert/bitwise_or/bitwise_right_shift/bitwise_xor

linalg namespace

(I copied those from the benchmark suite. If something turns out to be unfeasible we can drop it.)

• [ ] linalg.vecdot - finch.tensordot
• [ ] linalg.vector_norm
• [ ] linalg.trace
• [ ] linalg.tensordot - implemented in the main namespace. Just needs an alias
• [ ] linalg.outer
• [ ] linalg.cross
• [ ] linalg.matrix_transpose
• [ ] linalg.matrix_power
• [ ] linalg.matrix_norm - for nuc or 2, call external library. For fro, inf, 1, 0, -1, -inf, call jl.norm.
• [ ] xp.linalg.diagonal - finch.tensordot(finch.diagmask(), mtx)
• [ ] xp.linalg.cholesky - call CHOLMOD or something
• [ ] xp.linalg.det - call EIGEN or something
• [ ] xp.linalg.eigh - call external library
• [ ] xp.linalg.eigvalsh - call external library
• [x] xp.linalg.inv - call external library - scipy.sparse.linalg.inv
• [ ] xp.linalg.matrix_rank - call external library
• [ ] xp.linalg.pinv - call external library

Tensor methods and attributes

• [ ] Tensor.to_device() - finch.moveto

miscellaneous

• [ ] handling scalars
• [ ] benchmarking

[willow-ahrens]

I added some methods you can call for most of these! Some of them need to be external library calls

[mtsokol]

@willow-ahrens thank you - that's exactly what I needed! 😄