Is there a guide for adding additional deep learning framework support for CoLA ?

[hannibalhuang]

Per in the title

[mfinzi]

There's no guide at the moment, but which deep learning framework were you considering?

[hannibalhuang]

our team is developing the MindSpore deep learning framework, and would love to add support for CoLA :)

[mfinzi]

Ah well that would be great!
This week we are planning on doing some refactoring of jax_fns.py and torch_fns.py so as to conform more closely to how these functions are in the numpy namespace. After that we can write a guide or send over some pointers 👍

Right now the key points are

• a file providing each of the used operations in the same namespace convention
• extending AutoRegisteringPytree with the pytree registration in MindSpore if this is the approach that you guys use for vmap and related operations
• adding an additional case to get_library_fns
• adding an additional case to get_backend
• adding an additional parametrize case to the tests, e.g. here
• adding mindspore to the CI, probably we want this not in the matrix strategy but handled on its own

Probably we can streamline and simplify these requirements a little bit and will try to do so this week.

[mfinzi]

Working on this in #62

[mfinzi]

@hannibalhuang #62 is now merged to master and is in v0.0.4.
We have done some simplification of what is required for adding a new backend.

The steps to add a new framework are now:

• add a mindspore_fns.py in https://github.com/wilson-labs/cola/tree/main/cola/backends with each of the relevant functions we use (in the numpy syntax, see e.g. jax_fns.py or torch_fns.py)
• Implement the get_library_fns, get_xnp, and pytree registration in https://github.com/wilson-labs/cola/blob/main/cola/backends/backends.py
• list it in https://github.com/wilson-labs/cola/blob/main/cola/backends/__init__.py
• Add mindspore installation and run with the mindspore mark to the CI in https://github.com/wilson-labs/cola/blob/main/.github/workflows/run_test_suite.yml