LLama convert in TensorNetwork

Hello, I’m trying to convert a LLama model from safetonsors to TensorNetwork to reduce weight and speed up inference. I’m trying to integrate tensor layers into an existing model. However, I encounter errors. I want to retain all the knowledge of the model without training it. Has anyone tried this already? maybe there is a laptop of some kind so that we can see the working implementations

ValueError                                Traceback (most recent call last)
<ipython-input-34-c141f1326ef0> in <cell line: 36>()
     34 
     35 # Применение функций
---> 36 replace_with_tn_layers(model, tn_layers, test_input)
     37 
     38 

<ipython-input-34-c141f1326ef0> in replace_with_tn_layers(model, tn_layers, test_input)
     19             model._modules[name] = tn_layer
     20             if u.shape != tn_layer.kernel1.shape or v.shape != tn_layer.kernel2.shape:
---> 21                 raise ValueError(f"Несоответствие размеров для слоя {name}: размер u: {u.shape}, размер kernel1: {tn_layer.kernel1.shape}, размер v: {v.shape}, размер kernel2: {tn_layer.kernel2.shape}")
     22             tn_layer.kernel1.assign(tf.convert_to_tensor(u, dtype=tf.float32))
     23             tn_layer.kernel2.assign(tf.convert_to_tensor(v, dtype=tf.float32))

ValueError: Несоответствие размеров для слоя layers.0.self_attn.k_proj: размер u: (256, 64), размер kernel1: (2048, 64), размер v: (64, 2048), размер kernel2: (64, 256)

You are encountering a dimension mismatch error while integrating tensor network layers into an existing LLama model. The error arises because the shapes of the tensors u and v do not match the expected dimensions of kernel1 and kernel2 in the tensor network layers. To resolve this, you need to:

  1. Verify the dimensions of the original layers in the LLama model.
  2. Adjust the tensor network layers to match the dimensions of the corresponding LLama model layers.
  3. Consider reshaping tensors u and v if necessary, but be cautious as this might affect the model’s performance.
  4. Ensure the tensor network layers preserve the functionality of the original layers.
  5. Test and debug the model after adjustments to confirm successful integration.

There might not be a direct example or notebook for this specific case, so exploring TensorFlow and tensor network community forums and GitHub repositories could be beneficial for additional guidance.