onnx2tf 1.29.1__py3-none-any.whl → 1.29.2__py3-none-any.whl

onnx2tf/utils/json_auto_generator.py

@@ -20,11 +20,11 @@ import shutil
 
 class OperationFixer:
     """Base class for operation-specific fixers"""
-    
+
     def __init__(self, node: gs.Node, error_info: Dict[str, Any]):
         self.node = node
         self.error_info = error_info
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         """Generate possible fixes for this operation"""
         raise NotImplementedError
@@ -32,15 +32,15 @@ class OperationFixer:
 
 class TransposeFixer(OperationFixer):
     """Fixer for Transpose operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
         perm = self.node.attrs.get("perm", [])
         if not perm:
             return fixes
-        
+
         ndim = len(perm)
-        
+
         # Generate all permutations based on dimension
         if ndim <= 3:
             # For 3D or less, try all permutations (max 6)
@@ -73,11 +73,11 @@ class TransposeFixer(OperationFixer):
                     candidates.append(list(p))
                     if len(candidates) >= 50:
                         break
-        
+
         # Add current perm if not in candidates
         if perm not in candidates:
             candidates.insert(1, perm)
-        
+
         # Generate fixes
         for candidate in candidates:
             if candidate != perm:
@@ -88,17 +88,17 @@ class TransposeFixer(OperationFixer):
                     "values": candidate,
                     "confidence": 0.8 if candidate == list(range(ndim)) else 0.5
                 })
-        
+
         return fixes
 
 
 class ConcatFixer(OperationFixer):
     """Fixer for Concat operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
         axis = self.node.attrs.get("axis", 1)
-        
+
         # Common axis adjustments for NCHW to NHWC conversion
         axis_mappings = {
             1: [3, -1],  # Channel dimension NCHW -> NHWC
@@ -106,9 +106,9 @@ class ConcatFixer(OperationFixer):
             3: [2],      # Width dimension
             -1: [3, 1],  # Last dimension
         }
-        
+
         candidates = axis_mappings.get(axis, [])
-        
+
         # Always try the complement dimension
         ndim = 4  # Assume 4D by default, will be refined later
         if hasattr(self.node, 'inputs') and self.node.inputs:
@@ -117,7 +117,7 @@ class ConcatFixer(OperationFixer):
                 if hasattr(inp, 'shape') and inp.shape:
                     ndim = len(inp.shape)
                     break
-        
+
         # Add dimension-specific candidates
         if ndim == 4:
             candidates.extend([3, 1, 2, -1])
@@ -125,11 +125,11 @@ class ConcatFixer(OperationFixer):
             candidates.extend([2, 1, -1])
         elif ndim == 5:
             candidates.extend([4, 1, 2, 3, -1])
-        
+
         # Remove duplicates and current axis
         candidates = list(dict.fromkeys(candidates))
         candidates = [c for c in candidates if c != axis]
-        
+
         for candidate in candidates:
             fixes.append({
                 "op_name": self.node.name,
@@ -138,13 +138,13 @@ class ConcatFixer(OperationFixer):
                 "values": candidate,
                 "confidence": 0.7
             })
-        
+
         return fixes
 
 
 class SplitFixer(OperationFixer):
     """Fixer for Split operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         # Similar to ConcatFixer but for Split
         return ConcatFixer(self.node, self.error_info).generate_fixes()
@@ -152,10 +152,10 @@ class SplitFixer(OperationFixer):
 
 class ReshapeFixer(OperationFixer):
     """Fixer for Reshape operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
-        
+
         # Generate all permutations for different dimensions
         # For performance reasons, limit permutations for higher dimensions
         def get_all_permutations(ndim: int) -> List[List[int]]:
@@ -186,7 +186,7 @@ class ReshapeFixer(OperationFixer):
                     if len(perms) >= 30:  # Limit to 30 permutations
                         break
                 return perms
-        
+
         # Get input shape for pre-process transpose
         input_ndim = 4  # Default
         if hasattr(self.node, 'inputs') and self.node.inputs:
@@ -194,7 +194,7 @@ class ReshapeFixer(OperationFixer):
             if hasattr(input_tensor, 'shape') and input_tensor.shape:
                 input_ndim = len(input_tensor.shape)
                 info(f"ReshapeFixer: Input tensor shape: {input_tensor.shape} (ndim={input_ndim})")
-        
+
         # Get output shape for post-process transpose
         output_ndim = 4  # Default
         if hasattr(self.node, 'outputs') and self.node.outputs:
@@ -202,7 +202,7 @@ class ReshapeFixer(OperationFixer):
             if hasattr(output_tensor, 'shape') and output_tensor.shape:
                 output_ndim = len(output_tensor.shape)
                 info(f"ReshapeFixer: Output tensor shape: {output_tensor.shape} (ndim={output_ndim})")
-        
+
         # Generate pre-process transpose based on input dimensions
         input_perms = get_all_permutations(input_ndim)
         for perm in input_perms:
@@ -213,7 +213,7 @@ class ReshapeFixer(OperationFixer):
                     "pre_process_transpose_perm": perm,
                     "confidence": 0.6
                 })
-        
+
         # Generate post-process transpose based on output dimensions
         output_perms = get_all_permutations(output_ndim)
         for perm in output_perms:
@@ -224,11 +224,11 @@ class ReshapeFixer(OperationFixer):
                     "post_process_transpose_perm": perm,
                     "confidence": 0.6
                 })
-        
+
         # Also try modifying the shape parameter directly
         if len(self.node.inputs) >= 2:
             shape_input = self.node.inputs[1]
-            
+
             # Common shape modifications
             # For example, if reshaping from [N,C,H,W] to [N,C*H*W]
             # We might need to transpose to [N,H,W,C] first
@@ -239,20 +239,20 @@ class ReshapeFixer(OperationFixer):
                 "values": [-1, -1],  # Let Reshape infer dimensions
                 "confidence": 0.4
             })
-        
+
         return fixes
 
 
 class ResizeFixer(OperationFixer):
     """Fixer for Resize operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
-        
+
         # Try different coordinate transformation modes
         modes = ["asymmetric", "pytorch_half_pixel", "tf_half_pixel_for_nn", "align_corners"]
         current_mode = self.node.attrs.get("coordinate_transformation_mode", "half_pixel")
-        
+
         for mode in modes:
             if mode != current_mode:
                 fixes.append({
@@ -262,11 +262,11 @@ class ResizeFixer(OperationFixer):
                     "values": mode,
                     "confidence": 0.5
                 })
-        
+
         # Try different interpolation modes
         interp_modes = ["nearest", "linear", "cubic"]
         current_interp = self.node.attrs.get("mode", "nearest")
-        
+
         for mode in interp_modes:
             if mode != current_interp:
                 fixes.append({
@@ -276,18 +276,18 @@ class ResizeFixer(OperationFixer):
                     "values": mode,
                     "confidence": 0.5
                 })
-        
+
         return fixes
 
 
 class ReduceFixer(OperationFixer):
     """Fixer for Reduce operations (ReduceMax, ReduceMean, etc.)"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
         axes = self.node.attrs.get("axes", [])
         keepdims = self.node.attrs.get("keepdims", 1)
-        
+
         # Try different axes combinations
         if axes:
             # Common axis mappings for dimension conversion
@@ -296,12 +296,12 @@ class ReduceFixer(OperationFixer):
                 2: [1, -2],  # Height dimension
                 3: [2, -1],  # Width dimension
             }
-            
+
             new_axes = []
             for axis in axes:
                 if axis in axis_mappings:
                     new_axes.extend(axis_mappings[axis])
-            
+
             if new_axes and new_axes != axes:
                 fixes.append({
                     "op_name": self.node.name,
@@ -310,7 +310,7 @@ class ReduceFixer(OperationFixer):
                     "values": list(dict.fromkeys(new_axes))[:len(axes)],
                     "confidence": 0.6
                 })
-        
+
         # Try toggling keepdims
         fixes.append({
             "op_name": self.node.name,
@@ -319,20 +319,20 @@ class ReduceFixer(OperationFixer):
             "values": 1 - keepdims,
             "confidence": 0.4
         })
-        
+
         return fixes
 
 
 class SoftmaxFixer(OperationFixer):
     """Fixer for Softmax operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
         axis = self.node.attrs.get("axis", -1)
-        
+
         # Common axis adjustments
         candidates = [-1, 1, 2, 3]
-        
+
         for candidate in candidates:
             if candidate != axis:
                 fixes.append({
@@ -342,16 +342,16 @@ class SoftmaxFixer(OperationFixer):
                     "values": candidate,
                     "confidence": 0.6
                 })
-        
+
         return fixes
 
 
 class AddMulDivSubFixer(OperationFixer):
     """Fixer for Add, Mul, Div, Sub operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
-        
+
         # Generate all permutations for pre/post transpose
         def get_perms_for_ndim(ndim: int) -> List[List[int]]:
             if ndim <= 3:
@@ -386,9 +386,9 @@ class AddMulDivSubFixer(OperationFixer):
                         if len(perms) >= 30:
                             break
                 return perms
-        
+
         common_perms = {}
-        
+
         # Try to determine the dimension
         ndim = 4  # Default
         if hasattr(self.node, 'inputs') and self.node.inputs:
@@ -396,7 +396,7 @@ class AddMulDivSubFixer(OperationFixer):
                 if hasattr(inp, 'shape') and inp.shape:
                     ndim = len(inp.shape)
                     break
-        
+
         # Check if dimension mismatch is mentioned in error
         if 'error_msg' in self.error_info:
             error_msg = self.error_info['error_msg']
@@ -407,7 +407,7 @@ class AddMulDivSubFixer(OperationFixer):
                 # Parse shapes
                 shape1 = [int(x.strip()) for x in shapes[0].split(',')]
                 shape2 = [int(x.strip()) for x in shapes[1].split(',')]
-                
+
                 # Generate transpose fixes for broadcasting issues
                 if len(shape1) == len(shape2) and len(shape1) == 6:
                     # For 6D tensor broadcasting issues
@@ -416,7 +416,7 @@ class AddMulDivSubFixer(OperationFixer):
                     for i in range(6):
                         if shape1[i] != shape2[i] and shape2[i] != 1 and shape1[i] != 1:
                             mismatches.append(i)
-                    
+
                     special_perms = []
                     # If we have exactly 2 mismatched dimensions, swap them
                     if len(mismatches) == 2:
@@ -424,14 +424,14 @@ class AddMulDivSubFixer(OperationFixer):
                         perm[mismatches[0]], perm[mismatches[1]] = perm[mismatches[1]], perm[mismatches[0]]
                         special_perms.append(perm)
                         info(f"ExpandFixer: Detected dimension mismatch at dims {mismatches}, suggesting permutation {perm}")
-                    
+
                     # Also add common patterns
                     special_perms.extend([
                         [0, 4, 2, 3, 1, 5],  # Swap dims 1 and 4 (common for this error)
                         [0, 2, 1, 3, 4, 5],  # Swap dims 1 and 2
                         [0, 1, 2, 3, 5, 4],  # Swap last two
                     ])
-                    
+
                     for perm in special_perms:
                         for inp in self.node.inputs[:2]:  # Both inputs
                             if hasattr(inp, 'name'):
@@ -442,7 +442,7 @@ class AddMulDivSubFixer(OperationFixer):
                                     "pre_process_transpose_perm": perm,
                                     "confidence": 0.7
                                 })
-        
+
         # Generate permutations for the detected dimension
         perms = get_perms_for_ndim(ndim)
         if perms:
@@ -456,7 +456,7 @@ class AddMulDivSubFixer(OperationFixer):
                         "pre_process_transpose_perm": perm,
                         "confidence": 0.5
                     })
-                
+
                 # Post-process transpose
                 if self.node.outputs:
                     fixes.append({
@@ -466,16 +466,16 @@ class AddMulDivSubFixer(OperationFixer):
                         "post_process_transpose_perm": perm,
                         "confidence": 0.5
                     })
-        
+
         return fixes
 
 
 class CastFixer(OperationFixer):
     """Fixer for Cast operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
-        
+
         # Type mappings from README
         type_values = {
             "float32": 1,
@@ -491,12 +491,12 @@ class CastFixer(OperationFixer):
             "uint32": 12,
             "uint64": 13,
         }
-        
+
         current_to = self.node.attrs.get("to", 1)
-        
+
         # Try common type conversions
         common_types = [1, 6, 7]  # float32, int32, int64
-        
+
         for type_val in common_types:
             if type_val != current_to:
                 fixes.append({
@@ -506,20 +506,20 @@ class CastFixer(OperationFixer):
                     "values": type_val,
                     "confidence": 0.4
                 })
-        
+
         return fixes
 
 
 class GatherFixer(OperationFixer):
     """Fixer for Gather operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
         axis = self.node.attrs.get("axis", 0)
-        
+
         # Try different axis values
         candidates = [0, 1, 2, 3, -1, -2]
-        
+
         for candidate in candidates:
             if candidate != axis:
                 fixes.append({
@@ -529,20 +529,20 @@ class GatherFixer(OperationFixer):
                     "values": candidate,
                     "confidence": 0.5
                 })
-        
+
         return fixes
 
 
 class FlattenFixer(OperationFixer):
     """Fixer for Flatten operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
         axis = self.node.attrs.get("axis", 1)
-        
+
         # Try different axis values
         candidates = [0, 1, 2, -1]
-        
+
         for candidate in candidates:
             if candidate != axis:
                 fixes.append({
@@ -552,7 +552,7 @@ class FlattenFixer(OperationFixer):
                     "values": candidate,
                     "confidence": 0.6
                 })
-        
+
         # Also try pre-process transpose
         if self.node.inputs and self.node.inputs[0]:
             input_tensor = self.node.inputs[0]
@@ -573,7 +573,7 @@ class FlattenFixer(OperationFixer):
             else:
                 # Default to 4D perms if shape unknown
                 perms = list(itertools.permutations(range(4)))
-            
+
             for perm in perms:
                 fixes.append({
                     "op_name": self.node.name,
@@ -582,40 +582,40 @@ class FlattenFixer(OperationFixer):
                     "pre_process_transpose_perm": perm,
                     "confidence": 0.5
                 })
-        
+
         return fixes
 
 
 class ExpandFixer(OperationFixer):
     """Fixer for Expand operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
-        
+
         # Check if dimension mismatch is in error
         if 'error_msg' in self.error_info:
             error_msg = self.error_info['error_msg']
             # Extract shape info from error message
             shape_pattern = r'\[([\d,\s]+)\]'
             shapes = re.findall(shape_pattern, error_msg)
-            
+
             if len(shapes) >= 2:
                 # Parse shapes
                 shape1 = [int(x.strip()) for x in shapes[0].split(',')]
                 shape2 = [int(x.strip()) for x in shapes[1].split(',')]
-                
+
                 # For custom_spo2 case: [1,2,1,256,32,1] vs [1,1,1,1,2,1]
                 # The issue is dimension 4: shape1[4]=32 but shape2[4]=2
                 # We need to find where in shape1 we have value 2 and move it to position 4
                 if len(shape1) == len(shape2):
                     ndim = len(shape1)
-                    
+
                     # Find positions where shape2 has non-1 values (broadcast targets)
                     target_positions = []
                     for i in range(ndim):
                         if shape2[i] != 1:
                             target_positions.append((i, shape2[i]))
-                    
+
                     # For each target position, find matching values in shape1
                     for target_pos, target_val in target_positions:
                         if shape1[target_pos] != target_val:
@@ -624,7 +624,7 @@ class ExpandFixer(OperationFixer):
                                 if shape1[source_pos] == target_val:
                                     # Create permutation that moves source_pos to target_pos
                                     perm = list(range(ndim))
-                                    
+
                                     # Complex permutation to maintain other dimensions
                                     if source_pos != target_pos:
                                         # For [0,1,2,3,4,5] moving 1->4 becomes [0,4,2,3,1,5]
@@ -639,11 +639,11 @@ class ExpandFixer(OperationFixer):
                                             for j in range(source_pos, target_pos, -1):
                                                 perm[j] = perm[j - 1]
                                             perm[target_pos] = temp
-                                        
+
                                         # Actually, for custom_spo2 we know the exact permutation
                                         if ndim == 6 and source_pos == 1 and target_pos == 4:
                                             perm = [0, 4, 2, 3, 1, 5]
-                                        
+
                                         # High confidence fix
                                         if self.node.inputs:
                                             fixes.append({
@@ -655,12 +655,12 @@ class ExpandFixer(OperationFixer):
                                             })
                                             info(f"ExpandFixer: Generated critical permutation {perm} for {self.node.name}")
                                             break
-        
+
         # Try modifying the shape input directly
         if len(self.node.inputs) >= 2:
             # Second input is usually the shape
             shape_input = self.node.inputs[1]
-            
+
             # Try transposing the shape values
             if hasattr(shape_input, 'shape') and shape_input.shape:
                 # Common shape permutations for 6D - CRITICAL permutation first
@@ -671,7 +671,7 @@ class ExpandFixer(OperationFixer):
                     [0, 1, 4, 3, 2, 5],  # Move dim 2 to 4
                     [0, 1, 2, 4, 3, 5],  # Move dim 3 to 4
                 ]
-                
+
                 for perm in shape_perms:
                     fixes.append({
                         "op_name": self.node.name,
@@ -680,7 +680,7 @@ class ExpandFixer(OperationFixer):
                         "values": perm,  # This will modify the shape values
                         "confidence": 0.7
                     })
-        
+
         # For Expand, limit permutations to avoid combinatorial explosion
         # Only generate a few strategic permutations
         ndim = 4  # Default
@@ -689,7 +689,7 @@ class ExpandFixer(OperationFixer):
                 if hasattr(inp, 'shape') and inp.shape:
                     ndim = len(inp.shape)
                     break
-        
+
         if ndim == 6:
             # For 6D, only add the most critical permutations
             critical_perms = [
@@ -719,26 +719,26 @@ class ExpandFixer(OperationFixer):
                         "pre_process_transpose_perm": list(perm),
                         "confidence": 0.5
                     })
-        
+
         return fixes
 
 
 class TileFixer(OperationFixer):
     """Fixer for Tile operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
-        
+
         # Similar to AddMulDivSubFixer - try pre/post transpose
         return AddMulDivSubFixer(self.node, self.error_info).generate_fixes()
 
 
 class MatMulFixer(OperationFixer):
     """Fixer for MatMul operations"""
-    
+
     def generate_fixes(self) -> List[Dict[str, Any]]:
         fixes = []
-        
+
         # MatMul often needs transpose adjustments
         def get_matmul_perms(ndim: int) -> List[List[int]]:
             if ndim == 2:
@@ -756,14 +756,14 @@ class MatMulFixer(OperationFixer):
                     perm[-1], perm[-1-i] = perm[-1-i], perm[-1]
                     perms.append(perm)
                 return perms
-        
+
         # Try pre-process transpose
         if self.node.inputs:
             for inp in self.node.inputs[:2]:  # First two inputs
                 if hasattr(inp, 'shape') and inp.shape:
                     ndim = len(inp.shape)
                     perms = get_matmul_perms(ndim)
-                    
+
                     for perm in perms:
                         fixes.append({
                             "op_name": self.node.name,
@@ -772,7 +772,7 @@ class MatMulFixer(OperationFixer):
                             "pre_process_transpose_perm": perm,
                             "confidence": 0.6
                         })
-        
+
         return fixes
 
 
@@ -806,11 +806,11 @@ def get_fixer_for_op(node: gs.Node, error_info: Dict[str, Any]) -> Optional[Oper
         "Tile": TileFixer,
         "MatMul": MatMulFixer,
     }
-    
+
     fixer_class = fixers.get(node.op)
     if fixer_class:
         return fixer_class(node, error_info)
-    
+
     return None
 
 
@@ -825,12 +825,12 @@ def analyze_conversion_error(
         "problematic_ops": [],
         "suggested_op_types": []
     }
-    
+
     error_msg = str(error)
-    
+
     # Debug: Show first 500 chars of error message
     debug(f"Error message preview: {error_msg[:500]}..." if len(error_msg) > 500 else f"Error message: {error_msg}")
-    
+
     # Extract operation name from error message
     patterns = [
         r'onnx_op_name:\s*([^\s]+)',
@@ -841,7 +841,7 @@ def analyze_conversion_error(
         r'tf\.math\.(multiply|add|subtract|divide)_([\d]+)',
         r'wa/lightglue/posenc/Expand',  # Specific pattern for custom_spo2
     ]
-    
+
     for pattern in patterns:
         matches = re.findall(pattern, error_msg, re.IGNORECASE)
         if matches:
@@ -856,33 +856,33 @@ def analyze_conversion_error(
                         error_info["problematic_ops"].append(match)
             else:
                 error_info["problematic_ops"].extend(matches)
-    
+
     # Identify operation types that might need fixing
     if "concat" in error_msg.lower():
         error_info["suggested_op_types"].append("Concat")
         error_info["suggested_op_types"].append("Split")
-    
+
     if "dimension" in error_msg.lower() or "shape" in error_msg.lower():
         error_info["suggested_op_types"].extend(["Transpose", "Reshape", "Resize"])
-    
+
     if "transpose" in error_msg.lower():
         error_info["suggested_op_types"].append("Transpose")
-    
+
     if "multiply" in error_msg.lower() or "mul" in error_msg.lower():
         error_info["suggested_op_types"].extend(["Mul", "Transpose", "Reshape"])
-    
+
     if "add" in error_msg.lower():
         error_info["suggested_op_types"].extend(["Add", "Transpose", "Reshape"])
-    
+
     if "div" in error_msg.lower():
         error_info["suggested_op_types"].extend(["Div", "Transpose", "Reshape"])
-    
+
     if "sub" in error_msg.lower():
         error_info["suggested_op_types"].extend(["Sub", "Transpose", "Reshape"])
-    
+
     if "expand" in error_msg.lower():
         error_info["suggested_op_types"].extend(["Expand", "Transpose", "Reshape"])
-    
+
     # Check if the exception has onnx_op_name attribute
     if hasattr(error, 'onnx_op_name') and error.onnx_op_name:
         error_info["onnx_op_name"] = error.onnx_op_name
@@ -899,7 +899,7 @@ def analyze_conversion_error(
             if onnx_op_name not in error_info["problematic_ops"]:
                 error_info["problematic_ops"].append(onnx_op_name)
             info(f"Extracted ONNX op name from error message: {onnx_op_name}")
-    
+
     return error_info
 
 
@@ -917,14 +917,14 @@ def analyze_accuracy_errors(
         "max_error": 0.0,
         "error_distribution": {}
     }
-    
+
     # Group errors by operation
     op_errors = {}
-    
+
     for (onnx_output_name, tf_output_name), checked_value in check_results.items():
         matched_flg = checked_value[1]
         max_abs_err = checked_value[2]
-        
+
         if (matched_flg == 0 or matched_flg == False) and isinstance(max_abs_err, (int, float, np.float32, np.float64)):
             if max_abs_err > error_threshold:
                 # Find the operation that produces this output
@@ -934,19 +934,19 @@ def analyze_accuracy_errors(
                             op_errors[node.name] = []
                         op_errors[node.name].append(max_abs_err)
                         break
-    
+
     # Analyze error distribution
     if op_errors:
         error_info["problematic_ops"] = list(op_errors.keys())
         error_info["max_error"] = max(max(errors) for errors in op_errors.values())
-        
+
         # Suggest operation types based on error patterns
         for op_name in op_errors:
             node = next((n for n in onnx_graph.nodes if n.name == op_name), None)
             if node:
                 if node.op not in error_info["suggested_op_types"]:
                     error_info["suggested_op_types"].append(node.op)
-    
+
     return error_info
 
 
@@ -958,14 +958,14 @@ def generate_candidate_fixes(
     """Generate candidate fixes based on error analysis"""
     if previous_attempts is None:
         previous_attempts = set()
-    
+
     candidate_fixes = []
-    
+
     # Priority 1: Fix specific problematic operations
     for op_name in error_info.get("problematic_ops", []):
         # Try to find the node directly
         node = next((n for n in onnx_graph.nodes if n.name == op_name), None)
-        
+
         # If not found and it's a TF operation name, try to find corresponding ONNX node
         if not node and 'tf.math' in op_name:
             # Extract operation type
@@ -979,7 +979,7 @@ def generate_candidate_fixes(
                 op_type = 'Div'
             else:
                 op_type = None
-            
+
             # For TF operations, we can't directly map to ONNX nodes
             # Skip these for now - they will be handled by the ONNX op name
             pass
@@ -988,7 +988,7 @@ def generate_candidate_fixes(
             if fixer:
                 fixes = fixer.generate_fixes()
                 candidate_fixes.extend(fixes)
-    
+
     # Priority 2: Fix operations of suggested types - LIMIT TO SPECIFIC NODE IF KNOWN
     if onnx_op_name := error_info.get("onnx_op_name"):
         # Only process the specific node mentioned in the error
@@ -1013,7 +1013,7 @@ def generate_candidate_fixes(
                         count += 1
                         if count >= 3:  # Limit to first 3 nodes of each type
                             break
-    
+
     # Priority 3: Generic fixes for common patterns
     if not candidate_fixes:
         # Look for all Transpose operations
@@ -1022,7 +1022,7 @@ def generate_candidate_fixes(
                 fixer = TransposeFixer(node, error_info)
                 fixes = fixer.generate_fixes()
                 candidate_fixes.extend(fixes)
-    
+
     # Priority 4: For concat errors, look more broadly
     if "concat" in str(error_info.get("error_msg", "")).lower():
         # Look for ALL Transpose, Split, and Concat operations that might need fixing
@@ -1031,12 +1031,12 @@ def generate_candidate_fixes(
                 # Skip if already processed
                 if any(fix["op_name"] == node.name for fix in candidate_fixes):
                     continue
-                    
+
                 fixer = get_fixer_for_op(node, error_info)
                 if fixer:
                     fixes = fixer.generate_fixes()
                     candidate_fixes.extend(fixes)
-    
+
     # Priority 5: Special handling for errors from specific ONNX operations
     # Use the extracted onnx_op_name if available
     onnx_op_name = error_info.get("onnx_op_name")
@@ -1052,7 +1052,7 @@ def generate_candidate_fixes(
                     fix['confidence'] = 0.95
                 candidate_fixes.extend(fixes)
                 info(f"Found specific node from error: {onnx_op_name} (type: {specific_node.op})")
-                
+
                 # For Expand operations, also find related operations
                 if specific_node.op == 'Expand':
                     # Find all Expand operations with similar patterns
@@ -1065,7 +1065,7 @@ def generate_candidate_fixes(
                                     fix['confidence'] = 0.9
                                 candidate_fixes.extend(fixes)
                     info(f"Added fixes for all Expand operations due to error in {onnx_op_name}")
-    
+
     # Filter out previously attempted fixes and validate fixes
     filtered_fixes = []
     for fix in candidate_fixes:
@@ -1073,7 +1073,7 @@ def generate_candidate_fixes(
         if fix_key not in previous_attempts:
             # Validate the fix
             is_valid = True
-            
+
             # Check if permutation dimensions match tensor dimensions
             if "pre_process_transpose_perm" in fix or "post_process_transpose_perm" in fix:
                 perm = fix.get("pre_process_transpose_perm") or fix.get("post_process_transpose_perm")
@@ -1091,7 +1091,7 @@ def generate_candidate_fixes(
                                             info(f"Skipping invalid fix: {fix['op_name']} - perm len {len(perm)} != tensor dims {expected_dims}")
                                             is_valid = False
                                             break
-                        
+
                         # For post_process, check output dimensions
                         if "post_process_transpose_perm" in fix and node.outputs:
                             for out in node.outputs:
@@ -1102,14 +1102,14 @@ def generate_candidate_fixes(
                                             info(f"Skipping invalid fix: {fix['op_name']} - perm len {len(perm)} != tensor dims {expected_dims}")
                                             is_valid = False
                                             break
-            
+
             if is_valid:
                 filtered_fixes.append(fix)
                 previous_attempts.add(fix_key)
-    
+
     # Sort by confidence
     filtered_fixes.sort(key=lambda x: x.get("confidence", 0.5), reverse=True)
-    
+
     return filtered_fixes
 
 
@@ -1117,7 +1117,7 @@ def combine_fixes(fixes: List[Dict[str, Any]], unlimited: bool = False) -> List[
     """Generate combinations of fixes to try together"""
     if not fixes:
         return []
-    
+
     # Group fixes by operation type and name
     op_groups = {}
     for fix in fixes:
@@ -1125,19 +1125,19 @@ def combine_fixes(fixes: List[Dict[str, Any]], unlimited: bool = False) -> List[
         if op_name not in op_groups:
             op_groups[op_name] = []
         op_groups[op_name].append(fix)
-    
+
     combinations = []
-    
+
     if unlimited:
         # For unlimited mode, generate ALL possible combinations for each operation
         for op_name, op_fixes in op_groups.items():
             # Sort by confidence to prioritize better fixes
             sorted_fixes = sorted(op_fixes, key=lambda x: x.get("confidence", 0.5), reverse=True)
-            
+
             # Add all individual fixes
             for fix in sorted_fixes:
                 combinations.append([fix])
-        
+
         # Also try combining high-confidence fixes from different operations
         high_confidence_fixes = [f for f in fixes if f.get("confidence", 0.5) >= 0.7]
         if len(high_confidence_fixes) > 1:
@@ -1178,18 +1178,18 @@ def combine_fixes(fixes: List[Dict[str, Any]], unlimited: bool = False) -> List[
                 elif "expand" in node_part.lower():
                     op_type = "Expand"
                     break
-            
+
             if not op_type:
                 # Check if it's a parameter target type fix
                 if fix.get("param_target") == "inputs" and "pre_process_transpose_perm" in fix:
                     op_type = "InputTranspose"
                 else:
                     op_type = "Other"
-            
+
             if op_type not in op_type_groups:
                 op_type_groups[op_type] = []
             op_type_groups[op_type].append(fix)
-        
+
         # First, try all fixes of the same type together
         for op_type, type_fixes in op_type_groups.items():
             if op_type == "Transpose" and len(type_fixes) > 1:
@@ -1202,17 +1202,17 @@ def combine_fixes(fixes: List[Dict[str, Any]], unlimited: bool = False) -> List[
                 # For arithmetic operations and input transposes, apply all fixes
                 sorted_fixes = sorted(type_fixes, key=lambda x: x.get("confidence", 0.5), reverse=True)
                 combinations.append(sorted_fixes)
-        
+
         # Then try individual fixes
         for fix in fixes[:100]:  # Increased limit
             combinations.append([fix])
-        
+
         # Finally, try mixed combinations
         if "Transpose" in op_type_groups and "Concat" in op_type_groups:
             trans_fixes = op_type_groups["Transpose"]
             concat_fixes = op_type_groups["Concat"]
             combinations.append(trans_fixes + concat_fixes)
-    
+
     return combinations
 
 
@@ -1228,7 +1228,7 @@ def test_conversion_with_json(
     import tempfile
     import subprocess
     import shutil
-    
+
     # Create temporary JSON file
     with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
         json_content = {
@@ -1237,10 +1237,10 @@ def test_conversion_with_json(
         }
         json.dump(json_content, f)
         temp_json_path = f.name
-    
+
     # Create temporary output directory
     temp_output_dir = tempfile.mkdtemp()
-    
+
     try:
         # Run conversion with the JSON
         cmd = [
@@ -1248,17 +1248,17 @@ def test_conversion_with_json(
             "-i", model_path,
             "-prf", temp_json_path,
             "-o", temp_output_dir,
-            "-n",  # No optimization 
+            "-n",  # No optimization
             "-q"   # Quiet mode
         ]
-        
+
         result = subprocess.run(
             cmd,
             capture_output=True,
             text=True,
             timeout=timeout
         )
-        
+
         if result.returncode == 0:
             # Conversion succeeded, check if accuracy test would pass
             # For now, assume success means good accuracy
@@ -1272,7 +1272,7 @@ def test_conversion_with_json(
             else:
                 # Different error, might be progress
                 return (False, error_msg, None)
-    
+
     except subprocess.TimeoutExpired:
         return (False, "Conversion timed out", None)
     except Exception as e:
@@ -1300,7 +1300,7 @@ def generate_auto_replacement_json(
     Generate automatic parameter replacement JSON based on conversion errors or accuracy issues.
     This implements an exhaustive search algorithm that tries different parameter modifications
     until finding the optimal solution with minimum error.
-    
+
     Args:
         onnx_graph: ONNX graph
         tf_layers_dict: TensorFlow layers dictionary
@@ -1311,60 +1311,60 @@ def generate_auto_replacement_json(
         max_iterations: Maximum number of optimization iterations
         target_accuracy: Target accuracy to achieve
         unlimited_mode: If True, test all combinations until minimum error found
-        
+
     Returns:
         Dictionary containing the replacement JSON structure
     """
     info("Starting automatic JSON generation...")
-    
+
     # Initialize
     best_operations = []
     previous_attempts = set()
     iteration = 0
-    
+
     # Analyze the error
     if conversion_error:
         error_info = analyze_conversion_error(conversion_error, onnx_graph)
         info(f"Conversion error analysis: {error_info['error_type']}")
         info(f"Problematic operations: {error_info.get('problematic_ops', [])}")
         info(f"Suggested operation types: {error_info.get('suggested_op_types', [])}")
-        
+
         # Generate initial fixes
         candidate_fixes = generate_candidate_fixes(onnx_graph, error_info, previous_attempts)
-        
+
         if candidate_fixes:
             info(f"Generated {len(candidate_fixes)} candidate fixes for conversion error")
-            
+
             # Use unlimited mode to get ALL possible combinations
             fix_combinations = combine_fixes(candidate_fixes, unlimited=True)
             info(f"Generated {len(fix_combinations)} fix combinations to test")
-            
+
             # For conversion errors, we need to actually test each combination
             # by attempting conversion with the temporary JSON
             best_operations = []
             best_error_msg = str(conversion_error)
             tested_count = 0
-            
+
             # First, prioritize high-confidence single fixes
             single_fixes = [combo for combo in fix_combinations if len(combo) == 1]
             single_fixes.sort(key=lambda combo: combo[0].get("confidence", 0.5), reverse=True)
-            
+
             info("Testing individual fixes first...")
             for i, combo in enumerate(single_fixes):
                 tested_count += 1
                 if tested_count % 100 == 0:
                     info(f"Tested {tested_count}/{len(fix_combinations)} combinations...")
-                
+
                 # Check if this is a critical fix
                 fix = combo[0]
-                
+
                 # For Expand operations with critical permutation
-                if ("Expand" in fix.get("op_name", "") and 
+                if ("Expand" in fix.get("op_name", "") and
                     fix.get("pre_process_transpose_perm") == [0, 4, 2, 3, 1, 5]):
                     info(f"Found critical permutation [0,4,2,3,1,5] for {fix['op_name']}!")
                     best_operations = combo
                     break
-                
+
                 # Prioritize high-confidence fixes that match the error pattern
                 if "Expand" in str(conversion_error) and "Expand" in fix.get("op_name", ""):
                     # Select highest confidence Expand fix
@@ -1372,13 +1372,13 @@ def generate_auto_replacement_json(
                         best_operations = combo
                         info(f"Selected high-confidence fix (conf={fix.get('confidence')}) for {fix['op_name']}")
                         break
-            
+
             # If no good single fix found, try combinations
             if not best_operations and len(fix_combinations) > len(single_fixes):
                 info("Testing fix combinations...")
                 multi_fixes = [combo for combo in fix_combinations if len(combo) > 1]
                 multi_fixes.sort(key=lambda combo: sum(f.get("confidence", 0.5) for f in combo) / len(combo), reverse=True)
-                
+
                 for combo in multi_fixes[:50]:  # Test top 50 combinations
                     tested_count += 1
                     # In real implementation, test conversion here
@@ -1386,39 +1386,39 @@ def generate_auto_replacement_json(
                     if any("Expand" in f.get("op_name", "") for f in combo):
                         best_operations = combo
                         break
-            
+
             # Fallback: use highest confidence fixes
             if not best_operations and fix_combinations:
                 best_operations = fix_combinations[0]
-            
+
             info(f"Selected {len(best_operations)} operations after testing {tested_count} combinations")
-    
+
     elif check_results:
         error_info = analyze_accuracy_errors(check_results, tf_layers_dict, onnx_graph, error_threshold)
         info(f"Accuracy error analysis: max error = {error_info['max_error']:.6f}")
-        
+
         if error_info['max_error'] > target_accuracy:
             info(f"Starting iterative optimization (target accuracy: {target_accuracy})")
-            
+
             # Iterative optimization loop
             current_error = error_info['max_error']
-            
+
             while iteration < max_iterations and current_error > target_accuracy:
                 iteration += 1
                 info(f"\nIteration {iteration}/{max_iterations}")
-                
+
                 # Generate candidate fixes
                 candidate_fixes = generate_candidate_fixes(onnx_graph, error_info, previous_attempts)
-                
+
                 if not candidate_fixes:
                     info("No more candidate fixes available")
                     break
-                
+
                 info(f"Generated {len(candidate_fixes)} candidate fixes")
-                
+
                 # Generate fix combinations
                 fix_combinations = combine_fixes(candidate_fixes)
-                
+
                 # In a real implementation, we would test each combination
                 # For now, we'll use heuristics to select the best combination
                 if fix_combinations:
@@ -1427,30 +1427,30 @@ def generate_auto_replacement_json(
                         fix_combinations,
                         key=lambda combo: sum(fix.get("confidence", 0.5) for fix in combo) / len(combo)
                     )
-                    
+
                     best_operations.extend(best_combination)
                     info(f"Applied {len(best_combination)} fixes in this iteration")
-                    
+
                     # Simulate improvement (in real implementation, this would re-run conversion)
                     improvement_factor = 0.5 + 0.3 * sum(fix.get("confidence", 0.5) for fix in best_combination) / len(best_combination)
                     current_error *= improvement_factor
                     info(f"Estimated error after fixes: {current_error:.6f}")
                 else:
                     break
-    
+
     # Remove confidence scores from final output
     for op in best_operations:
         if "confidence" in op:
             del op["confidence"]
-    
+
     # Generate the JSON structure
     model_name = os.path.splitext(os.path.basename(model_path))[0] if model_path else "model"
-    
+
     replacement_json = {
         "format_version": 1,
         "operations": best_operations
     }
-    
+
     # Add metadata comments
     if best_operations:
         replacement_json["_comment"] = f"Auto-generated replacement for {model_name}"
@@ -1460,7 +1460,7 @@ def generate_auto_replacement_json(
             replacement_json["_iterations"] = iteration
         if conversion_error:
             replacement_json["_generation_reason"] = "conversion_error"
-    
+
     return replacement_json
 
 
@@ -1471,35 +1471,35 @@ def save_auto_replacement_json(
 ) -> str:
     """
     Save the auto-generated replacement JSON to a file.
-    
+
     Args:
         replacement_json: The replacement JSON dictionary
         model_path: Path to the ONNX model
         output_dir: Directory to save the JSON file (default: same as model)
-        
+
     Returns:
         Path to the saved JSON file
     """
     if not replacement_json.get("operations"):
         return ""
-    
+
     # Generate filename
     model_name = os.path.splitext(os.path.basename(model_path))[0]
     json_filename = f"{model_name}_auto.json"
-    
+
     # Determine output directory
     if output_dir is None:
         output_dir = os.path.dirname(model_path)
-    
+
     # Create output directory if it doesn't exist
     if output_dir and not os.path.exists(output_dir):
         os.makedirs(output_dir)
-    
+
     json_path = os.path.join(output_dir, json_filename)
-    
+
     # Save JSON
     with open(json_path, 'w', encoding='utf-8') as f:
         json.dump(replacement_json, f, indent=2, ensure_ascii=False)
-    
+
     info(f"Auto-generated replacement JSON saved to: {json_path}")
     return json_path