Assess kernel completeness

[yanndupis]

In the last few months, we have added lots of operation. We would like to assess the completeness for each operations in terms of types we support.

[TODO] give more detail to this issue.

The assessment can be found in this file: https://docs.google.com/spreadsheets/d/1mTYJdGzofBoWSzUAWF0zqHo8qSdy0qYagbTQaSh4RVU/edit#gid=0

[jvmncs]

We (@yanndupis and I) made a separate sheet for the logical layer only, to help assess user-facing functionality completeness. Some observations from that exercise:

• We added an asterisk on the Op row when we felt like it needed more discussion, or resolution of open issues (e.g. ReplicatedPublic vs Mirrored)
• We tried to maintain parity across a few columns. Generally, we want all Fixed@Replicated ops to be available on Host as well, so that a user can move a fixedpoint computation from Replicated to Host while debugging it.
• We also want all Fixedpoint@Host ops to have Float@Host analogues (because Fixed should be a subset of Float)
• Going through the exercise, it's not obvious that this is the final breakdown of all the Ops we should have. Maybe we should be stricter about the API we are using as a target (e.g. exact Numpy API or TF API whenever possible, any exceptions strictly documented). This deserves a deeper discussion, and these discussions are prob higher priority than filling in most of these green boxes.
• We discussed the value of having complete coverage over all dtypes (following e.g. numpy). Decided that maybe this might be valuable eventually but not for Phase 1/2 (prob not even Phase 3) of OSS release
• Discussed the same for type promotion semantics (which is what pytorch/numpy/jax/tf-numpy use, but tf does not). Decided that Moose should not have any type promotion semantics for now, but e.g. pymoose (or a specific integration with any of the above) could handle that on its own if desired, as long as we have casting fully implemented between the types we support

[jvmncs]

We also made a sheet called slim-logical. The idea behind this was to find the minimal set of kernels we feel need to be tackled to support parity across Fixed@Repl, Fixed@Host, and Float@Host just for the PyMoose predictors we have today. There are a few exceptions for kernels we felt were critical for any array-based language, mostly related to indexing and reshaping.

[jvmncs]

We also took a look at jax.lax, the XLA primitives library in Jax, to get a sense of how close we are to being able to support a full-featured ML library with this list. We chose that one because it's significantly smaller than both TF and PyTorch, but still big enough to be interesting.