Skip to content

tenzr::ArgSortNode

More...

#include <node.h>

Inherits from tenzr::Node, std::enable_shared_from_this< Node >

Public Functions

Name
ArgSortNode(std::shared_ptr< Node > values, std::shared_ptr< Node > indices, bool descending =false, int64_t axis =0, bool return_values =false)
virtual std::string OpName() const override
StableHLO-style name of this node's operation.
const std::shared_ptr< Node > & Values() const
const std::shared_ptr< Node > & Indices() const
bool Descending() const
Whether the sort is descending.
int64_t Axis() const
Axis along which to sort.
bool ReturnValues() const
Whether to return sorted values instead of indices.

Additional inherited members

Public Functions inherited from tenzr::Node

Name
Node(NodeKind kind, Shape shape, DType dtype, ShardingSpec sharding =ShardingSpec::Replicated())
virtual ~Node() =default
NodeKind Kind() const
The node's operation kind.
const Shape & shape() const
The node's output shape.
DType Dtype() const
The node's output dtype.
const ShardingSpec & Sharding() const
The node's sharding spec.
void SetSharding(ShardingSpec spec)
Replace the sharding spec.
const StridedLayout & Layout() const
The node's strided layout.
void SetLayout(StridedLayout layout)
Replace the layout.
const std::vector< std::shared_ptr< Node > > & Inputs() const
The node's input operands.
uint64_t Id() const
The node's unique id.

Protected Functions inherited from tenzr::Node

Name
uint64_t NextId()

Protected Attributes inherited from tenzr::Node

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

Detailed Description

class tenzr::ArgSortNode;

Sort an operand by value along an axis, returning either indices or sorted values.

Public Functions Documentation

function ArgSortNode

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

Parameters:

  • values Values to sort.
  • indices Index tensor carried alongside the values.
  • descending Sort order.
  • axis Sort axis.
  • return_values Return sorted values when true, indices otherwise.

Create an argsort node from values and their indices.

function OpName

inline virtual std::string OpName() const override

StableHLO-style name of this node's operation.

Reimplements: tenzr::Node::OpName

function Values

inline const std::shared_ptr< Node > & Values() const

function Indices

inline const std::shared_ptr< Node > & Indices() const

function Descending

inline bool Descending() const

Whether the sort is descending.

function Axis

inline int64_t Axis() const

Axis along which to sort.

function ReturnValues

inline bool ReturnValues() const

Whether to return sorted values instead of indices.