owlbarn / owl

Owl - OCaml Scientific Computing @ https://ocaml.xyz
MIT License
1.22k stars 124 forks source link

`set_fancy` does not work for some arrays with certain `Bigarray.kind` types. #671

Open zoj613 opened 5 months ago

zoj613 commented 5 months ago

When using Owl.Dense.Ndarray.Generic.set_fancy with arrays of kinds not Float32, Float64, Complex32 or Complex64, the operation fails with error message Exception: Failure "_ndarray_set_fancy: unsupported operation", even though the kind is a valid Bigarray.kind type.

It appears this exception is thrown by the snippet: https://github.com/owlbarn/owl/blob/600c21a233a6d3c17675c219e4d9f6397a4b7a0f/src/owl/core/owl_slicing_fancy.ml#L92-L105

I think it is very limiting to only support just 4 kinds when the set operation should work for any of the Bigarray.kind types.

Here is a minimal example to reproduce the error:

module Ndarray = Owl.Dense.Ndarray.Generic
let arr = Ndarray.create Bigarray.Int32 [|3;4;3|] 5l
let slice = Owl_types.[L [1; 2]; R []; I 0]
let x = Ndarray.create Bigarray.Int32 [|2; 4; 1|] 0l
 Ndarray.set_fancy slice arr x

I expected this simple operation to work but it throws a |Exception: Failure "_ndarray_set_fancy: unsupported operation". exception.

I would love for Owl to support more kinds as im using it to write an array library that needs to support more than the 4 basic kinds listed above.

Cc @mseri @jzstark

zoj613 commented 4 months ago

I was able to implement a temporary workaround to this issue by using element-wise set. Basically:

  1. Convert the slice to a list of coordinates of the target array.
  2. Use Owl.Dense.Ndarray.Generic.set and List.iter2 to iterate over the list of coordinates and values and set them one at a time
 List.iter2
  (fun coord val ->
    Owl.Dense.Ndarray.Generic.set coord val) coord_list value_list

I would like to add that this issue also affects the transpose function as well. The offending lines are: https://github.com/owlbarn/owl/blob/455bde1ac6f1f94c24eca43bdbe45af4f80e139f/src/owl/core/owl_ndarray_transpose.ml#L41-L54

Unfortunately, I haven't been able to find a workaround for this

zoj613 commented 4 months ago

UPDATE: So I managed to find a workaround for computing the transpose of any kind using Owl functions. Essentially, this involves converting between .Any and .Generic representations of an array. The functions that seem to not work on .Generic happen to work on .Any arrays. Here is an implementation of transpose that works for any Bigarray kind:

module Ndarray = Owl.Dense.Ndarray.Generic
module Anyarray = Owl.Dense.Ndarray.Any

let transpose ?axis x =
  let shape = Ndarray.shape x in
  let y = Anyarray.init_nd shape @@ fun c -> Ndarray.get x c in
  let y' = Anyarray.transpose ?axis y in
  Ndarray.init_nd (Ndarray.kind x) (Anyarray.shape y') @@ fun c ->
    Anyarray.get y' c

I suppose this workaround can be generalized to any of the Generic functions that are limited to only 4 bigarray kinds?

Cc @jzstark

mseri commented 4 months ago

Nice. The main downside is that this won't scale well since it involves copying everything twice, but at least would circumvent the current unnecessary limitation

zoj613 commented 4 months ago

Nice. The main downside is that this won't scale well since it involves copying everything twice, but at least would circumvent the current unnecessary limitation

Perhaps it can be used only for the Bigarray kinds not currently supported by the Generic module, that's if correctness/usability is prioritized over performance.

I am very interested in getting Generic functions to work for any kind. What part of the codebase would I need to extend to support this functionality? Is it doable or too complicated for someone who doesn't have intimate knowledge of the codebase?

jzstark commented 2 months ago

Thanks a lot the investigation about using other types of Bigarray! Unfortunately, due to that the type system holds a central position in the design of Owl, adding even one new type such as Float16 or Int requires daunting work which is not available to us now.