Saved Model support: implements Model.saved_model that enables the user to serialize pred_step as a SavedModel

[cgarciae]

saved_model gets a sample input and a path, and serializes a pred_step as a SavedModel via jax2tf:

import numpy as np
import elegy
from flax import linen

x = np.random.uniform(size=(32, 6))

model = elegy.Model(module=linen.Dense(4))
model.init(x)

model.saved_model(x, "/saved/model/path", batch_size=[1, 8])

This creates a SavedModel with 2 signatures:

• input: (1, 6)
• output: (1, 4)

and

• input: (8, 6)
• output: (8, 4)

Due to google/jax#5915 only a finite set of static batch sizes work, if this is resolved in the future one will be able to omit the batch size to que a polymorphic batch dimension:

import numpy as np
import elegy
from flax import linen

x = np.random.uniform(size=(32, 6))

model = elegy.Model(module=linen.Dense(4))
model.init(x)

model.saved_model(x, "/saved/model/path")

This SavedModel would contain the following signature:

• input: (None, 6)
• output: (None, 4)

[codecov-io]

Codecov Report

Merging #171 (fb484cc) into master (7f4a781) will increase coverage by 0.01%.
The diff coverage is 88.88%.

@@            Coverage Diff             @@
##           master     #171      +/-   ##
==========================================
+ Coverage   87.21%   87.22%   +0.01%     
==========================================
  Files         135      136       +1     
  Lines        7360     7430      +70     
==========================================
+ Hits         6419     6481      +62     
- Misses        941      949       +8     

Impacted Files	Coverage Δ
elegy/model/utils.py	82.14% <82.14%> (ø)
elegy/model/model_core.py	85.46% <86.95%> (+0.23%)	⬆️
elegy/model/model_test.py	100.00% <100.00%> (ø)
elegy/nn/linear.py	97.29% <100.00%> (+0.07%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 7f4a781...fb484cc. Read the comment docs.