ml_switcheroo.plugins.mlx_optimizers
====================================

.. py:module:: ml_switcheroo.plugins.mlx_optimizers

.. autoapi-nested-parse::

   Plugin for MLX Optimizer translation.

   Handles impedance mismatches between PyTorch and Apple MLX optimizers.



Functions
---------

.. autoapisummary::

   ml_switcheroo.plugins.mlx_optimizers.transform_mlx_optimizer_init
   ml_switcheroo.plugins.mlx_optimizers.transform_mlx_optimizer_step
   ml_switcheroo.plugins.mlx_optimizers.transform_mlx_zero_grad


Module Contents
---------------

.. py:function:: transform_mlx_optimizer_init(node: libcst.Call, ctx: ml_switcheroo.core.hooks.HookContext) -> libcst.Call

.. py:function:: transform_mlx_optimizer_step(node: libcst.Call, ctx: ml_switcheroo.core.hooks.HookContext) -> Union[libcst.Call, libcst.FlattenSentinel]

.. py:function:: transform_mlx_zero_grad(node: libcst.Call, ctx: ml_switcheroo.core.hooks.HookContext) -> libcst.CSTNode