| 67 | |
| 68 | # Transformer to replace Dense with MyDense |
| 69 | class LayerReplacer(ast.NodeTransformer): |
| 70 | def visit_Call(self, node): |
| 71 | if isinstance(node.func, ast.Name) and node.func.id == 'Dense': |
| 72 | node.func.id = custom_layer |
| 73 | return self.generic_visit(node) |
| 74 | |
| 75 | transformer = LayerReplacer() |
| 76 | modified_tree = transformer.visit(tree) |