node.h¶

Namespaces¶

Name
tenzr
Classes¶

	Name
class	tenzr::Node
class	tenzr::VariableNode
class	tenzr::AddNode Elementwise addition of two operands.
class	tenzr::MulNode Elementwise multiplication of two operands.
class	tenzr::SubNode Elementwise subtraction of two operands.
class	tenzr::DivNode Elementwise division of two operands.
class	tenzr::MatmulNode Matrix multiplication of two operands.
class	tenzr::ConvolutionNode
class	tenzr::ReshapeNode View of an operand reshaped to a new shape.
class	tenzr::TransposeNode View of an operand with its dimensions reordered by a permutation.
class	tenzr::SliceNode View of an operand restricted by per-axis slice bounds.
class	tenzr::PadNode Pad an operand with a fill value at the edges and between elements.
class	tenzr::GatherNode Gather slices of an operand at the given indices.
class	tenzr::ScatterNode
class	tenzr::DynamicSliceNode
class	tenzr::DynamicUpdateSliceNode
class	tenzr::ArgSortNode
class	tenzr::ConcatenateNode Join multiple operands along an axis.
class	tenzr::ReluNode Elementwise rectified linear unit.
class	tenzr::SigmoidNode Elementwise logistic sigmoid.
class	tenzr::CustomOpNode Node backed by a custom operator registered in the OpRegistry.
class	tenzr::ConvertNode Convert an operand to another element type.
class	tenzr::ReduceNode
class	tenzr::ReduceWindowNode
Source code¶

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef TENZR_FRONTEND_NODE_H_
#define TENZR_FRONTEND_NODE_H_

#include <atomic>
#include <memory>
#include <string>
#include <vector>

#include "tenzr/frontend/dtype.h"
#include "tenzr/frontend/layout.h"
#include "tenzr/frontend/memory_pool.h"
#include "tenzr/frontend/shape.h"
#include "tenzr/frontend/sharding.h"
#include "tenzr/runtime/device_buffer.h"

namespace tenzr {

  struct OpDef;

  enum class NodeKind {
    kVariable,
    kAdd,
    kMul,
    kSub,
    kDiv,
    kMatmul,
    kConvolution,
    kReshape,
    kTranspose,
    kSlice,
    kRelu,
    kSigmoid,
    kCustomOp,
    kConvert,
    kReduce,
    kReduceWindow,
    kPad,
    kGather,
    kScatter,
    kDynamicSlice,
    kDynamicUpdateSlice,
    kArgSort,
    kConcatenate
  };

  class Node : public std::enable_shared_from_this<Node> {
   public:
    Node(NodeKind kind, Shape shape, DType dtype, ShardingSpec sharding = ShardingSpec::Replicated());
    virtual ~Node() = default;

    NodeKind Kind() const { return kind_; }

    const Shape& shape() const { return shape_; }

    DType Dtype() const { return dtype_; }

    const ShardingSpec& Sharding() const { return sharding_; }

    void SetSharding(ShardingSpec spec) { sharding_ = std::move(spec); }

    const StridedLayout& Layout() const { return layout_; }

    void SetLayout(StridedLayout layout) { layout_ = std::move(layout); }

    const std::vector<std::shared_ptr<Node>>& Inputs() const { return inputs_; }

    uint64_t Id() const { return id_; }

    virtual std::string OpName() const = 0;

   protected:
    NodeKind kind_;
    Shape shape_;
    DType dtype_;
    ShardingSpec sharding_;
    StridedLayout layout_;
    std::vector<std::shared_ptr<Node>> inputs_;
    uint64_t id_;

    static uint64_t NextId();
  };

  class VariableNode : public Node {
   public:
    VariableNode(Shape shape, DType dtype, std::string name = "", ShardingSpec sharding = ShardingSpec::Replicated());

    std::string OpName() const override { return "Variable"; }

    const std::string& VariableName() const { return name_; }

    void SetData(std::shared_ptr<PinnedBuffer> data);

    const std::shared_ptr<PinnedBuffer>& Data() const;

    bool HasData() const;

    void SetDeviceData(std::shared_ptr<DeviceBuffer> device_data);

    const std::shared_ptr<DeviceBuffer>& DeviceData() const;

    bool HasDeviceData() const;

    void SyncHost();

    void SyncDevice();

   private:
    std::string name_;
    mutable std::shared_ptr<PinnedBuffer> data_;
    std::shared_ptr<DeviceBuffer> device_data_;
  };

  class AddNode : public Node {
   public:
    AddNode(std::shared_ptr<Node> lhs, std::shared_ptr<Node> rhs);

    std::string OpName() const override { return "Add"; }

    const std::shared_ptr<Node>& Lhs() const { return inputs_[0]; }
    const std::shared_ptr<Node>& Rhs() const { return inputs_[1]; }
  };

  class MulNode : public Node {
   public:
    MulNode(std::shared_ptr<Node> lhs, std::shared_ptr<Node> rhs);

    std::string OpName() const override { return "Mul"; }

    const std::shared_ptr<Node>& Lhs() const { return inputs_[0]; }
    const std::shared_ptr<Node>& Rhs() const { return inputs_[1]; }
  };

  class SubNode : public Node {
   public:
    SubNode(std::shared_ptr<Node> lhs, std::shared_ptr<Node> rhs);

    std::string OpName() const override { return "Sub"; }

    const std::shared_ptr<Node>& Lhs() const { return inputs_[0]; }
    const std::shared_ptr<Node>& Rhs() const { return inputs_[1]; }
  };

  class DivNode : public Node {
   public:
    DivNode(std::shared_ptr<Node> lhs, std::shared_ptr<Node> rhs);

    std::string OpName() const override { return "Div"; }

    const std::shared_ptr<Node>& Lhs() const { return inputs_[0]; }
    const std::shared_ptr<Node>& Rhs() const { return inputs_[1]; }
  };

  class MatmulNode : public Node {
   public:
    MatmulNode(std::shared_ptr<Node> lhs, std::shared_ptr<Node> rhs);

    std::string OpName() const override { return "Matmul"; }

    const std::shared_ptr<Node>& Lhs() const { return inputs_[0]; }
    const std::shared_ptr<Node>& Rhs() const { return inputs_[1]; }
  };

  class ConvolutionNode : public Node {
   public:
    ConvolutionNode(std::shared_ptr<Node> input, std::shared_ptr<Node> kernel, std::vector<int64_t> window_strides, std::vector<int64_t> lhs_dilation,
                    std::vector<int64_t> rhs_dilation, std::vector<int64_t> padding_low, std::vector<int64_t> padding_high);

    std::string OpName() const override { return "Convolution"; }

    const std::shared_ptr<Node>& Input() const { return inputs_[0]; }

    const std::shared_ptr<Node>& Kernel() const { return inputs_[1]; }

    const std::vector<int64_t>& WindowStrides() const { return window_strides_; }

    const std::vector<int64_t>& LhsDilation() const { return lhs_dilation_; }

    const std::vector<int64_t>& RhsDilation() const { return rhs_dilation_; }

    const std::vector<int64_t>& PaddingLow() const { return padding_low_; }

    const std::vector<int64_t>& PaddingHigh() const { return padding_high_; }

   private:
    std::vector<int64_t> window_strides_;
    std::vector<int64_t> lhs_dilation_;
    std::vector<int64_t> rhs_dilation_;
    std::vector<int64_t> padding_low_;
    std::vector<int64_t> padding_high_;
  };

  class ReshapeNode : public Node {
   public:
    ReshapeNode(std::shared_ptr<Node> input, Shape new_shape);

    std::string OpName() const override { return "Reshape"; }

    const std::shared_ptr<Node>& Input() const { return inputs_[0]; }
  };

  class TransposeNode : public Node {
   public:
    TransposeNode(std::shared_ptr<Node> input, std::vector<int64_t> permutation);

    std::string OpName() const override { return "Transpose"; }

    const std::shared_ptr<Node>& Input() const { return inputs_[0]; }

    const std::vector<int64_t>& Permutation() const { return permutation_; }

   private:
    std::vector<int64_t> permutation_;
  };

  class SliceNode : public Node {
   public:
    SliceNode(std::shared_ptr<Node> input, std::vector<Slice> slices);

    std::string OpName() const override { return "Slice"; }

    const std::shared_ptr<Node>& Input() const { return inputs_[0]; }

    const std::vector<Slice>& Slices() const { return slices_; }

   private:
    std::vector<Slice> slices_;
  };

  class PadNode : public Node {
   public:
    PadNode(std::shared_ptr<Node> operand, std::shared_ptr<Node> padding_value, std::vector<int64_t> edge_padding_low, std::vector<int64_t> edge_padding_high,
            std::vector<int64_t> interior_padding);

    std::string OpName() const override { return "Pad"; }

    const std::shared_ptr<Node>& Operand() const { return inputs_[0]; }

    const std::shared_ptr<Node>& PaddingValue() const { return inputs_[1]; }

    const std::vector<int64_t>& EdgePaddingLow() const { return low_; }

    const std::vector<int64_t>& EdgePaddingHigh() const { return high_; }

    const std::vector<int64_t>& InteriorPadding() const { return interior_; }

   private:
    std::vector<int64_t> low_;
    std::vector<int64_t> high_;
    std::vector<int64_t> interior_;
  };

  class GatherNode : public Node {
   public:
    GatherNode(std::shared_ptr<Node> operand, std::shared_ptr<Node> indices);

    std::string OpName() const override { return "Gather"; }

    const std::shared_ptr<Node>& Operand() const { return inputs_[0]; }
    const std::shared_ptr<Node>& Indices() const { return inputs_[1]; }
  };

  class ScatterNode : public Node {
   public:
    ScatterNode(std::shared_ptr<Node> operand, std::shared_ptr<Node> indices, std::shared_ptr<Node> updates, std::string update_op = "add");

    std::string OpName() const override { return update_op_ == "add" ? "ScatterAdd" : "Scatter"; }

    const std::string& UpdateOp() const { return update_op_; }

    const std::shared_ptr<Node>& Operand() const { return inputs_[0]; }

    const std::shared_ptr<Node>& Indices() const { return inputs_[1]; }

    const std::shared_ptr<Node>& Updates() const { return inputs_[2]; }

   private:
    std::string update_op_;
  };

  class DynamicSliceNode : public Node {
   public:
    DynamicSliceNode(std::shared_ptr<Node> operand, std::vector<std::shared_ptr<Node>> start_indices, std::vector<int64_t> slice_sizes);

    std::string OpName() const override { return "DynamicSlice"; }

    const std::shared_ptr<Node>& Operand() const { return inputs_[0]; }

    const std::vector<int64_t>& SliceSizes() const { return slice_sizes_; }

   private:
    std::vector<int64_t> slice_sizes_;
  };

  class DynamicUpdateSliceNode : public Node {
   public:
    DynamicUpdateSliceNode(std::shared_ptr<Node> operand, std::shared_ptr<Node> update, std::vector<std::shared_ptr<Node>> start_indices);

    std::string OpName() const override { return "DynamicUpdateSlice"; }

    const std::shared_ptr<Node>& Operand() const { return inputs_[0]; }
    const std::shared_ptr<Node>& Update() const { return inputs_[1]; }
  };

  class ArgSortNode : public Node {
   public:
    ArgSortNode(std::shared_ptr<Node> values, std::shared_ptr<Node> indices, bool descending = false, int64_t axis = 0, bool return_values = false);

    std::string OpName() const override { return return_values_ ? "SortValues" : "ArgSort"; }

    const std::shared_ptr<Node>& Values() const { return inputs_[0]; }

    const std::shared_ptr<Node>& Indices() const { return inputs_[1]; }

    bool Descending() const { return descending_; }

    int64_t Axis() const { return axis_; }

    bool ReturnValues() const { return return_values_; }

   private:
    bool descending_{false};
    int64_t axis_{0};
    bool return_values_{false};
  };

  class ConcatenateNode : public Node {
   public:
    ConcatenateNode(std::vector<std::shared_ptr<Node>> inputs, int64_t axis);

    std::string OpName() const override { return "Concat"; }

    int64_t Axis() const { return axis_; }

   private:
    int64_t axis_{0};
  };

  class ReluNode : public Node {
   public:
    explicit ReluNode(std::shared_ptr<Node> input);
    std::string OpName() const override { return "Relu"; }

    const std::shared_ptr<Node>& Input() const { return inputs_[0]; }
  };

  class SigmoidNode : public Node {
   public:
    explicit SigmoidNode(std::shared_ptr<Node> input);

    std::string OpName() const override { return "Sigmoid"; }

    const std::shared_ptr<Node>& Input() const { return inputs_[0]; }
  };

  class CustomOpNode : public Node {
   public:
    CustomOpNode(std::string op_name, std::vector<std::shared_ptr<Node>> inputs, const OpDef* Definition = nullptr);

    std::string OpName() const override { return op_name_str_; }

    const OpDef* Definition() const { return op_def_; }

   private:
    std::string op_name_str_;
    const OpDef* op_def_{nullptr};
  };

  class ConvertNode : public Node {
   public:
    ConvertNode(std::shared_ptr<Node> input, DType target_dtype);

    std::string OpName() const override { return "Convert"; }

    const std::shared_ptr<Node>& Input() const { return inputs_[0]; }
  };

  enum class ReduceOp { kSum, kProd, kMax, kMin };

  class ReduceNode : public Node {
   public:
    ReduceNode(std::shared_ptr<Node> input, ReduceOp Reduction, std::vector<int64_t> dimensions, bool keepdims);

    std::string OpName() const override;

    const std::shared_ptr<Node>& Input() const { return inputs_[0]; }

    ReduceOp Reduction() const { return reduce_op_; }

    const std::vector<int64_t>& Dimensions() const { return dimensions_; }

    const std::vector<int64_t>& CanonicalDimensions() const { return canonical_dimensions_; }

    bool Keepdims() const { return keepdims_; }

    const Shape& NonKeepdimsShape() const { return non_keepdims_shape_; }

    Shape ExpandedShape() const;

   private:
    ReduceOp reduce_op_;
    std::vector<int64_t> dimensions_;
    std::vector<int64_t> canonical_dimensions_;
    bool keepdims_;
    Shape non_keepdims_shape_;
  };

  class ReduceWindowNode : public Node {
   public:
    ReduceWindowNode(std::shared_ptr<Node> input, std::shared_ptr<Node> init_value, std::vector<int64_t> window_dimensions, std::vector<int64_t> window_strides,
                     std::vector<int64_t> padding_low, std::vector<int64_t> padding_high, std::vector<int64_t> window_dilations, std::string Reduction);

    std::string OpName() const override { return "ReduceWindow"; }

    const std::shared_ptr<Node>& Input() const { return inputs_[0]; }

    const std::shared_ptr<Node>& InitValue() const { return inputs_[1]; }

    const std::vector<int64_t>& WindowDimensions() const { return window_dimensions_; }

    const std::vector<int64_t>& WindowStrides() const { return window_strides_; }

    const std::vector<int64_t>& PaddingLow() const { return padding_low_; }

    const std::vector<int64_t>& PaddingHigh() const { return padding_high_; }

    const std::vector<int64_t>& WindowDilations() const { return window_dilations_; }

    const std::string& Reduction() const { return reduce_op_; }

   private:
    std::vector<int64_t> window_dimensions_;
    std::vector<int64_t> window_strides_;
    std::vector<int64_t> padding_low_;
    std::vector<int64_t> padding_high_;
    std::vector<int64_t> window_dilations_;
    std::string reduce_op_;
  };

}  // namespace tenzr

#endif  // TENZR_FRONTEND_NODE_H_