# Ir2Torch and Torch2Ir

The `Ir2Torch` module is responsible for converting an intermediate representation (IR) back into a PyTorch model. This is useful for scenarios where you need to reconstruct a PyTorch model from its IR, which might have been generated by the `Torch2Ir` module.

## Concepts

### Intermediate Representation (IR)

The IR is a structured format that represents a PyTorch model in a way that is independent of the PyTorch framework. It captures the model's architecture, including layers, connections, and attributes.


## Extending Ir2Torch

To extend `Ir2Torch` with support for new modules, you need to register custom handlers that define how to convert specific IR nodes back into PyTorch modules.

### Step-by-Step Guide

1. **Define a Custom Linear Handler Function**
   Create a function that takes an `Implementation` object and returns a corresponding PyTorch module.

   ```python
   import torch.nn as nn
   from elasticai.creator.torch2ir import Implementation
   from elasticai.creator.ir2torch import get_default_converter as get_ir2torch

   my_ir2torch = get_ir2torch()

   @my_ir2torch.register_override("linear")
   def custom_linear_handler(impl: Implementation) -> nn.Module:
       return nn.Linear(
           in_features=impl.data["in_features"],
           out_features=impl.data["out_features"],
           bias=impl.data["bias"]
       )
   ```

2. **Convert the IR**
   Use the `Ir2Torch` instance to convert the IR back into a PyTorch model.

   ```python
   ir = {
       "root": {
           "name": "root",
           "type": "module",
           "nodes": {
               "input_1": {"name": "input_1", "type": "input", "implementation": "input"},
               "_0": {"name": "_0", "type": "custom_linear", "implementation": "0"},
               "output": {"name": "output", "type": "output", "implementation": "output"},
           },
           "edges": {
               ("input_1", "_0"): {"src": "input_1", "dst": "_0"},
               ("_0", "output"): {"src": "_0", "dst": "output"},
           },
       },
       "0": {
           "name": "0",
           "type": "custom_linear",
           "in_features": 1,
           "out_features": 2,
           "bias": True,
           "edges": {},
           "nodes": {},
       },
   }

   reconstructed_model = ir2torch_converter.convert(ir)
   ```

By following these steps, you can extend `Ir2Torch` to support custom modules and ensure that your IR can be accurately converted back into a PyTorch model.

:::{important}
The reconstructed model will not contain the state of an original
pytorch model. Often the generated module hierarchy will be compatible
with the state dict, that you can obtain by calling `torch.nn.Module.state_dict()`
and you can load that dict after rebuilding the model with `torch.nn.Module.load_state(state)`.
In some cases, e.g., if the Ir graph is altered or if the original module contained
and used more than a single reference to the same submodule, loading the state dict
directly will fail. In these cases you have to modify the state dict, such that
the paths to parameters from the original model match the paths to the parameters
in the rebuilt model.

```python
lin = nn.Linear(1, 1)
seq = Sequential(lin, lin)
```

In the example above the module `seq` contains two references to the same module.
When you read the state dict you will obtain a set of parameters for each of these
duplicate submodules with names like `"0.weight"` and `"1.weight"`.
During the generation of our Ir the duplicate instance will be ignored and the graph
will just contain two calls to the `lin` module.
Thus, you will have to remove the parameters starting with `"1."` from the state dict
before loading.
:::

### Extending Torch2Ir

Extending Torch2Ir works very similar, but with a few differences:

* There is currently no support for overriding handlers.
* You register new handlers with the `register` or `register_handlers` method.
