jax.export.export#
- jax.export.export(fun_jit, *, platforms=None, disabled_checks=())[源代码]#
导出 JAX 函数以进行持久序列化。
- 参数:
fun_jit (stages.Wrapped) – 要导出的函数。应为 jax.jit 的结果。
platforms (Sequence[str] | None | None) – 可选序列,包含“tpu”、“cpu”、“cuda”、“rocm”的子集。如果指定了多个平台,则导出的代码会接受一个指定平台的参数。如果为 None,则使用默认的 JAX 后端。有关多个平台的调用约定,请参见 https://jax.ac.cn/en/latest/export/export.html#module-calling-convention。
disabled_checks (Sequence[DisabledSafetyCheck]) – 要禁用的安全检查。请参阅 jax.export.DisabledSafetyCheck 的文档。
- 返回:
一个函数,它接受 {class}`jax.ShapeDtypeStruct` 的参数和关键字参数 pytrees,或具有 .shape 和 .dtype 属性的值,并返回一个 Exported。
- 返回类型:
Callable[…, Exported]
用法
>>> from jax import export >>> exported: export.Exported = export.export(jnp.sin)( ... np.arange(4, dtype=np.float32)) >>> >>> # You can inspect the Exported object >>> exported.in_avals (ShapedArray(float32[4]),) >>> blob: bytearray = exported.serialize() >>> >>> # The serialized bytes are safe to use in a separate process >>> rehydrated: export.Exported = export.deserialize(blob) >>> rehydrated.fun_name 'sin' >>> rehydrated.call(np.array([.1, .2, .3, .4], dtype=np.float32)) Array([0.09983342, 0.19866933, 0.29552022, 0.38941833], dtype=float32)