Implement naive `tree_flatten` and `tree_unflatten`.

ray-project / rayfed

A multiple parties joint, distributed execution engine based on Ray, to help build your own federated learning frameworks in minutes.

https://rayfed.readthedocs.io

Apache License 2.0

92 stars 22 forks source link

Implement naive `tree_flatten` and `tree_unflatten`. #111

Closed jovany-wang closed 1 year ago

jovany-wang commented 1 year ago

In this PR, we're introducing a naive implementation for tree_flatten and tree_unflatten API, for removing jax dependencies purpose.

After this PR, the RayFed doesn't depend on jax/jaxlib any longer. Note that, this naive implementation might be not as efficient as jax implementation.

close https://github.com/ray-project/rayfed/issues/106

zhouaihui commented 1 year ago

Nice feature! Only some mirror questions.

jovany-wang commented 1 year ago

@zhouaihui Sorry for that I didn't clean the code. Will ping you after it gets done.

jovany-wang commented 1 year ago

@zhouaihui @NKcqx I didn't fix. Please take a look, thanks~