Source code for afe.ir.serializer.visualizer_defines

#########################################################
# Copyright (C) 2023 SiMa Technologies, Inc.
#
# This material is SiMa proprietary and confidential.
#
# This material may not be copied or distributed without
# the express prior written permission of SiMa.
#
# All rights reserved.
#########################################################
# Code owner: Christopher Rodrigues
#########################################################
"""
Python analogues of data structures for the Netron file format.
These are temporary data structures representing data that
will be written to a JSON file for Netron.
"""
import dataclasses
from dataclasses import dataclass
from typing import NewType, List, Any, Dict

# Node name in Netron file.  Used as ID for references to nodes.
[docs] NNodeName = NewType("NNodeName", str)
@dataclass(frozen=True)
[docs] class NArray: """Properties of an array, to be represented in a Keras JSON file's weightsManifest."""
[docs] dtype: str
[docs] shape: List[int]
[docs] name: str
[docs] def to_json(self) -> Any: return dataclasses.asdict(self)
@dataclass(frozen=True)
[docs] class NLayerWeights: """Constant arrays associated with a layer, to be represented in a Keras JSON file's weightsManifest."""
[docs] weights: List[NArray]
[docs] def to_json(self) -> Any: # 'paths' is required, and we do not use it return {'paths': [], 'weights': [a.to_json() for a in self.weights]}
@dataclass(frozen=True)
[docs] class NInboundNode: """ Properties of an inbound node, to be represented in a Keras JSON file. An inbound node corresponds to one input graph edge of an AwesomeNode. """
[docs] name: NNodeName # Node that is passed as input
[docs] index: int # Index in this node's parameter list
[docs] def to_json(self) -> Any: # Two extra items are required (their meaning is unknown), and we do not use them return [self.name, 0, self.index, {}]
@dataclass(frozen=True)
[docs] class NLayer: """ Properties of a layer, to be represented in a Keras JSON file. A layer corresponds to an AwesomeNode. """
[docs] class_name: str # Displayed as graph node label
[docs] name: NNodeName # Displayed in attributes
[docs] inbound_nodes: List[NInboundNode] # Displayed as graph edges with details in attributes
[docs] config: Dict[str, Any] # Displayed as key-value pairs in attributes
[docs] def to_json(self) -> Any: return { 'class_name': self.class_name, 'name': self.name, 'inbound_nodes': [[n.to_json() for n in self.inbound_nodes]], # Extra list is intended 'config': self.config }
@dataclass(frozen=True)
[docs] class NInterfaceLayer: """ Properties of one input or output of a model, to be represented in a Keras JSON file. """
[docs] layer_name: NNodeName
[docs] layer_num: int
[docs] def to_json(self): return [self.layer_name, self.layer_num, self.layer_num] # Duplicate layer_num is intended