tenzr::Node¶

More...

#include <node.h>

Inherits from std::enable_shared_from_this< Node >

Inherited by tenzr::AddNode, tenzr::ArgSortNode, tenzr::ConcatenateNode, tenzr::ConvertNode, tenzr::ConvolutionNode, tenzr::CustomOpNode, tenzr::DivNode, tenzr::DynamicSliceNode, tenzr::DynamicUpdateSliceNode, tenzr::GatherNode, tenzr::MatmulNode, tenzr::MulNode, tenzr::PadNode, tenzr::ReduceNode, tenzr::ReduceWindowNode, tenzr::ReluNode, tenzr::ReshapeNode, tenzr::ScatterNode, tenzr::SigmoidNode, tenzr::SliceNode, tenzr::SubNode, tenzr::TransposeNode, tenzr::VariableNode

Public Functions¶

Protected Functions¶

Protected Attributes¶

Detailed Description¶

class tenzr::Node;

Base class for nodes in the lazy expression graph. Each concrete node carries its operation kind, output shape, dtype, sharding, layout, and input operands.

Public Functions Documentation¶

function Node¶

Node(
    NodeKind kind,
    Shape shape,
    DType dtype,
    ShardingSpec sharding =ShardingSpec::Replicated()
)

Parameters:

• kind Node kind.
• shape Output shape.
• dtype Output dtype.
• sharding Output sharding.

Initialize a node with its kind, shape, dtype, and sharding.

function ~Node¶

virtual ~Node() =default

function Kind¶

inline NodeKind Kind() const

The node's operation kind.

function shape¶

inline const Shape & shape() const

The node's output shape.

function Dtype¶

inline DType Dtype() const

The node's output dtype.

function Sharding¶

inline const ShardingSpec & Sharding() const

The node's sharding spec.

function SetSharding¶

inline void SetSharding(
    ShardingSpec spec
)

Replace the sharding spec.

function Layout¶

inline const StridedLayout & Layout() const

The node's strided layout.

function SetLayout¶

inline void SetLayout(
    StridedLayout layout
)

Replace the layout.

function Inputs¶

inline const std::vector< std::shared_ptr< Node > > & Inputs() const

The node's input operands.

function Id¶

inline uint64_t Id() const

The node's unique id.

function OpName¶

virtual std::string OpName() const =0

StableHLO-style name of this node's operation.

Reimplemented by: tenzr::VariableNode::OpName, tenzr::AddNode::OpName, tenzr::MulNode::OpName, tenzr::SubNode::OpName, tenzr::DivNode::OpName, tenzr::MatmulNode::OpName, tenzr::ConvolutionNode::OpName, tenzr::ReshapeNode::OpName, tenzr::TransposeNode::OpName, tenzr::SliceNode::OpName, tenzr::PadNode::OpName, tenzr::GatherNode::OpName, tenzr::ScatterNode::OpName, tenzr::DynamicSliceNode::OpName, tenzr::DynamicUpdateSliceNode::OpName, tenzr::ArgSortNode::OpName, tenzr::ConcatenateNode::OpName, tenzr::ReluNode::OpName, tenzr::SigmoidNode::OpName, tenzr::CustomOpNode::OpName, tenzr::ConvertNode::OpName, tenzr::ReduceNode::OpName, tenzr::ReduceWindowNode::OpName

Protected Functions Documentation¶

function NextId¶

static uint64_t NextId()

Return: The new id.

Allocate a fresh, process-unique node id.

Protected Attributes Documentation¶

variable kind_¶

NodeKind kind_;

variable shape_¶

Shape shape_;

variable dtype_¶

DType dtype_;

variable sharding_¶

ShardingSpec sharding_;

variable layout_¶

StridedLayout layout_;

variable inputs_¶

std::vector< std::shared_ptr< Node > > inputs_;

variable id_¶

uint64_t id_;