jax.experimental.serialize_executable module

Contents

jax.experimental.serialize_executable module#

Pickling support for precompiled binaries.

API#

serialize(compiled)

Serializes a compiled binary.

deserialize_and_load(serialized, in_tree, ...)

Constructs a jax.stages.Compiled from a serialized executable.