ml_switcheroo.compiler.backends.python_snippet
==============================================

.. py:module:: ml_switcheroo.compiler.backends.python_snippet

.. autoapi-nested-parse::

   Python Snippet Emitter.

   This module provides the `PythonSnippetEmitter`, a specialized backend component
   designed to generate atomic Python statements (CST nodes) from LogicalNodes.

   Unlike `PythonBackend` which synthesizes entire modules/classes, this emitter
   focuses on generating individual:
   1.  **Initialization Statements**: `self.layer = LayerClass(...)`
   2.  **Execution Statements**: `output = self.layer(input)`
   3.  **Expressions**: `self.layer(input)` or `func(input)`

   This is used by the ``GraphPatcher`` to surgically insert new code for
   nodes created during Graph Optimization (e.g. Fused Operations).



Classes
-------

.. autoapisummary::

   ml_switcheroo.compiler.backends.python_snippet.PythonSnippetEmitter


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

.. py:class:: PythonSnippetEmitter(framework: str = 'torch')

   Generates isolated Python statements from LogicalNodes.

   Adapts generation logic based on the target framework to ensure
   idiomatic code (e.g., injecting `rngs` for Flax, `nn.` prefix for Torch).


   .. py:attribute:: framework
      :value: 'torch'



   .. py:method:: emit_init(node: ml_switcheroo.compiler.ir.LogicalNode) -> libcst.SimpleStatementLine

      Generates the initialization statement for a stateful layer.

      :param node: The logical node describing the layer.

      :returns: A LibCST statement node representing `self.id = Kind(...)`.



   .. py:method:: emit_call(node: ml_switcheroo.compiler.ir.LogicalNode, input_vars: List[str], output_var: str) -> libcst.SimpleStatementLine

      Generates the execution call statement (Assignment).

      :param node: The logical node.
      :param input_vars: List of variable names to pass as arguments.
      :param output_var: Variable name to assign the result to.

      :returns: A LibCST statement node (Assign).



   .. py:method:: emit_expression(node: ml_switcheroo.compiler.ir.LogicalNode, input_vars: List[str]) -> libcst.BaseExpression

      Generates the function call expression (without assignment).

      :param node: The logical node.
      :param input_vars: List of input variable names.

      :returns: A LibCST expression node.