# Copyright (c) Meta Platforms, Inc. and affiliates.

# All rights reserved.

#

# This source code is licensed under the BSD-style license found in the

# LICENSE file in the root directory of this source tree.

# pyre-unsafe

"""

Base Intermediate Representation for Developer Tools consumers

(e.g. TensorBoard, Terminal Debugger)

"""

from __future__ import annotations

from dataclasses import dataclass

from typing import Any, Dict, List, Optional

# Base Representation of a generic node within a ModelGraph

@dataclass

class Node:

name: str

# Nodes that this Node consumes/in-edges

inputs: Optional[List[Node]] = None

# List of output shapes

output_shapes: Optional[List[List[int]]] = None

# Generic Node level metadata

metadata: Optional[Dict[str, Any]] = None

# Names of the arguments derived from the op schema:

named_args: Optional[List[str]] = None

# Base Representation of an operator subgraph with metadata

@dataclass

class OperatorGraph:

# Identifier used for grouping nodes (e.g. expand/minimize Module)

graph_name: str

# Nodes and Sub-Graphs

elements: List[Node | OperatorGraph]

# Graph Level Metadata

metadata: Optional[Dict[str, Any]] = None

"""

Node SubClasses Types

"""

# Representation of a "Value" node within a ModelGraph

# i.e. Non-Operator Nodes

@dataclass

class ValueNode(Node):

dtype: str = ""

val: Optional[Any] = None

# Representation of an "OP" node within a ModelGraph

@dataclass

class OperatorNode(Node):

op: Optional[str] = None