将 JAX 前端支持添加到 Ivy 翻译器
作者: YushaArif99创建于 2024年12月17日更新于 2025年5月5日
标签JAX FrontendToDoTranspiler
- Add Native Framework-Specific Implementations for Core Transformation Passes:
- For example, implement the
native_jax_recursive_transformer.pyfor traversing and transforming JAX native source code. - Use
native_torch_recursive_transformer.pyas a reference (example here)
- Define the Transformation Pipeline for JAX to JAX Frontend IR:
- Create a new pipeline in
source_to_frontend_translator_config.pyto handle the stagesource='jax', target='jax_frontend'(example here).
- Define the Transformation Pipeline for JAX Frontend IR to Ivy:
- Add another pipeline in
frontend_to_ivy_translator_config.pyto handle the stagesource='jax_frontend', target='ivy'(example here).
- Add Stateful Classes for Flax APIs:
- Implement a stateful class for the
flax.nnx.ModuleAPI that inherits fromivy.Module. - Reference the existing implementation for PyTorch's
nn.Module(example here) - This allows for sequential lowering:
内容来源: unifyai/ivy