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 #90
• 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
• Boolean indexing (https://data-apis.org/array-api/latest/API_specification/indexing.html#boolean-array-indexing)

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 #89
• 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.

[ilan-gold]

Hello!

Re: #88 I think the root of this is having a way to do boolean indexing for us.

I don't see anything here about boolean indexing (https://data-apis.org/array-api/latest/API_specification/indexing.html#boolean-array-indexing) along axes. And from what I can tell, it's not supported in the julia library (or it is but only as a custom loop: https://github.com/finch-tensor/Finch.jl/blob/2e89d2228a777a8238eccd145cfa6a7b405eb76c/docs/src/docs/language/mask_sugar.md?plain=1#L5); however it is part of the array-api.

Is there a plan to support it? I would really love to use this but lacking a proper indexing method is a big blocker (assuming I have this right).

[willow-ahrens]

@ilan-gold Hi Ilan! Thanks for your interest in the project! Our goal is to support all of the operations in the array API, so if something has been left out it likely wasn't intentional. It is quite feasible to implement this operation, and I've added it to the list! The only caveat is that I don't think we will be able to implement it in a way that can fuse it with other operations, at least for now. I think I left it out of the julia implementation just because I needed to get things up and running quickly, but Boolean indexing is actually part of the Julia indexing api as well, so this addition is in scope for both languages!

[ilan-gold]

that can fuse it with other operations

Could you highlight what this would mean practically for users? This is mostly a performance, not a practical (in the sense of "all downstream operations will break if you boolean index" thing), no?

[willow-ahrens]

Yes, this is a performance thing, not a semantics thing. What it would mean for users is just that if we wish to do a logical indexing operation, we would need to materialize the inputs and outputs of the operation (This is normal for most libraries, Finch just provides a special fusion optimization for certain operations).

[ilan-gold]

Right, ok, interesting. Thanks @willow-ahrens I'll keep my ear to the ground in the meantime. Your library is a big target for us using the array-api, really looking forward to digging into it a bit more!