Exception: Unimplemented nodes

[kabyanil]

I am converting the modules of a transformer model from pytorch to tflite. Here is my conversion code -

sample_input = torch.randint(low=0, high=60, size =(1, 32, 512), dtype=torch.float)

src_pos_module = model.src_pos.eval()

src_pos_keras = nobuco.pytorch_to_keras(
    src_pos_module,
    args=[sample_input], kwargs=None,
    inputs_channel_order=ChannelOrder.TENSORFLOW,
    outputs_channel_order=ChannelOrder.TENSORFLOW,
    debug_traces=True
)

src_pos_keras.save(f"{common_dir}/src_pos.keras")

Here is my implementation of the positional encoding layer -

class PositionalEncoding(nn.Module):

    def __init__(self, d_model: int, seq_len: int, dropout: float) -> None:
        super().__init__()
        self.d_model = d_model
        self.seq_len = seq_len
        self.dropout = nn.Dropout(dropout)
        # Create a matrix of shape (seq_len, d_model)
        pe = torch.zeros(seq_len, d_model)
        # Create a vector of shape (seq_len)
        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1) # (seq_len, 1)
        # Create a vector of shape (d_model)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) # (d_model / 2)
        # Apply sine to even indices
        pe[:, 0::2] = torch.sin(position * div_term) # sin(position * (10000 ** (2i / d_model))
        # Apply cosine to odd indices
        pe[:, 1::2] = torch.cos(position * div_term) # cos(position * (10000 ** (2i / d_model))
        # Add a batch dimension to the positional encoding
        pe = pe.unsqueeze(0) # (1, seq_len, d_model)
        # Register the positional encoding as a buffer
        self.register_buffer('pe', pe)

    def forward(self, x):

        # x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False) # (batch, seq_len, d_model)
        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False) # (batch, seq_len, d_model)
        return self.dropout(x)

Here is the error I am getting -

[Nobuco] Tracing (DONE): 4 ops [00:00]
[Nobuco] Converting (DONE): |██████▋   | 2/3 ops [00:00]
Legend:
    Green — conversion successful
    Yellow — conversion imprecise
    Red — conversion failed
    Red — no converter found
    Bold — conversion applied directly
    * — subgraph reused
    Tensor — this output is not dependent on any of subgraph's input tensors
    Tensor — this input is a parameter / constant
    Tensor — this tensor is useless

PositionalEncoding[__main__](float32_0<1,32,512>) -> float32_2<1,32,512>
 │  requires_grad_[torch.Tensor](float32_1<1,32,512>, False) -> float32_1<1,32,512>
 │  __add__[torch.Tensor](float32_0<1,32,512>, float32_1<1,32,512>) -> float32_2<1,32,512>
 │   └· add[TensorBase](float32_0<1,32,512>, float32_1<1,32,512>) -> float32_2<1,32,512>
 ├· Dropout[torch.nn.modules.dropout](float32_2<1,32,512>) -> float32_2<1,32,512>
 └   └· dropout[torch.nn.functional](float32_2<1,32,512>, 0.1, False, False) -> float32_2<1,32,512>

Unimplemented nodes:
PositionalEncoding[__main__](float32_0<1,32,512>) -> float32_2<1,32,512>
 │  requires_grad_[torch.Tensor](float32_1<1,32,512>, False) -> float32_1<1,32,512>

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/nobuco/converters/validation.py", line 47, in validate
    diffs = validate_diff_default(keras_op, pytorch_op, input_args, input_kwargs, output_tensors)
  File "/usr/local/lib/python3.10/dist-packages/nobuco/converters/validation.py", line 72, in validate_diff_default
    outputs_tf = keras_op(*args_tf, **kwargs_tf)
  File "/usr/local/lib/python3.10/dist-packages/nobuco/layers/stub.py", line 7, in __call__
    raise Exception(f'Unimplemented op: {self.original_node}')
Exception: Unimplemented op: <method 'requires_grad_' of 'torch._C.TensorBase' objects>
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/nobuco/converters/validation.py", line 47, in validate
    diffs = validate_diff_default(keras_op, pytorch_op, input_args, input_kwargs, output_tensors)
  File "/usr/local/lib/python3.10/dist-packages/nobuco/converters/validation.py", line 70, in validate_diff_default
    keras_op.reset_states()
  File "/usr/local/lib/python3.10/dist-packages/nobuco/layers/container.py", line 82, in reset_states
    op.reset_states()
AttributeError: 'UnimplementedOpStub' object has no attribute 'reset_states'
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
[<ipython-input-36-cd7ea28a6126>](https://localhost:8080/#) in <cell line: 7>()
      5 src_pos_module = model.src_pos.eval()
      6 
----> 7 src_pos_keras = nobuco.pytorch_to_keras(
      8     src_pos_module,
      9     args=[sample_input], kwargs=None,

[/usr/local/lib/python3.10/dist-packages/nobuco/convert.py](https://localhost:8080/#) in pytorch_to_keras(model, args, kwargs, input_shapes, inputs_channel_order, outputs_channel_order, input_names, output_names, trace_shape, enable_torch_tracing, constants_to_variables, full_validation, validation_tolerance, return_outputs_pt, save_trace_html, debug_traces)
    363         print('Unimplemented nodes:')
    364         print(unimplemented_hierarchy.__str__(**vis_params))
--> 365         raise Exception('Unimplemented nodes')
    366 
    367     keras_op = keras_converted_node.keras_op

Exception: Unimplemented nodes