google / google/orbax

Enable saving (and loading) zero size arrays

Open
#2,587 0 comments 0 reactions 0 assignees View on GitHub
type:feature
Dominant language
Python
Stars
535
Forks
101
Avg merge
1d 13h
Merged PRs (30d)
23

Description

Dear all,

I wonder if it would be possible to support saving and loading zero size arrays. I saw it was deactivated in https://github.com/google/orbax/pull/1570. Allowing it would be very useful for us for saving rank 0 lora adapters for https://github.com/NovaSky-AI/SkyRL/tree/main/skyrl-tx (we use rank 0 lora adapters to allow users to selectively deactivate adapters for certain layers). Currently we have to write a bunch of code to work around this, whereas I don't see a reason why orbax shouldn't be able to handle it, since zero size arrays are valid arrays.

All the best and let us know if you have any further questions about the use case.

-- Philipp.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.