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

• astype - API: finch.astype function #15 - eager
• element-wise ops (add, multiply, cos, ...) - API: Lazy API #17 (partially...)
• reductions (xp.prod, xp.sum) - jl.sum and jl.prod, also just jl.reduce - API: Lazy API #17
• matmul - implemented with finch.tensordot for non-stacked input. Should be rewritten with jl.mul / Finch einsum.
• tensordot - finch.tensordot - API: Implement tensordot and matmul #22
• where - jl.broadcast(jl.ifelse, cond, a, b) - API: Implement where and nonzero #30
• argmin/argmax - jl.argmin (bug willow if this isn't implemented already) - eager for now
• take - jl.getindex eager for now
• nonzero - this is an eager function, but it is implemented as ffindnz(arr) - API: Implement where and nonzero #30
• 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. API: Add asarray function #28, API: Add eye function #32
• stats functions: max, mean, min, std, var
• set functions: unique_all, unique_counts, unique_inverse, unique_values - eager
• all, any
• concat - eager for now
• expand_dims - lazy
• flip -eager for now
• reshape - eager for now
• roll - eager for now
• squeeze - lazy
• stack - eager for now
• argsort/sort - eager
• broadcast_arrays - eager for now
• broadcast_to - eager for now
• can_cast/finfo/iinfo/result_type
• 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 - finch.norm
• linalg.trace - eager
• linalg.tensordot - implemented in the main namespace. Just needs an alias
• linalg.outer
• linalg.cross - eager for now
• linalg.matrix_transpose - lazy
• linalg.matrix_power - eager (call matmul on sparse matrix until it gets too dense)
• 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
• 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
• Don't forget about array_namespace_info() function, and the isdtype() and result_type() data type functions 🙂

[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! 😄

[jorenham]

Don't forget about __array_namespace_info__() function, and the isdtype() and result_type() data type functions 🙂

[willow-ahrens]

Thanks! Added a bullet for it.