include/tenzr/plax/plax.h¶

// 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_PLAX_PLAX_H_
#define TENZR_PLAX_PLAX_H_

#include <string_view>
#include <vector>

#include "tenzr/array.h"
#include "tenzr/common/export.h"

namespace tenzr {
  class NamedSharding;
}  // namespace tenzr

namespace tenzr::plax {
  struct TopKResult {
    Array values;
    Array indices;
  };


  // Strict Elementwise Binary Primitives

  [[nodiscard]] TENZR_API Array Add(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Sub(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Mul(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Div(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Pow(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Maximum(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Minimum(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Remainder(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Atan2(const Array& lhs, const Array& rhs);

  // Strict Comparison Primitives (returns Bool array)

  [[nodiscard]] TENZR_API Array Eq(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Ne(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Lt(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Le(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Gt(const Array& lhs, const Array& rhs);
  [[nodiscard]] TENZR_API Array Ge(const Array& lhs, const Array& rhs);

  // Strict Elementwise Unary Primitives

  [[nodiscard]] TENZR_API Array Neg(const Array& x);
  [[nodiscard]] TENZR_API Array Exp(const Array& x);
  [[nodiscard]] TENZR_API Array Log(const Array& x);
  [[nodiscard]] TENZR_API Array tanh(const Array& x);
  [[nodiscard]] TENZR_API Array Abs(const Array& x);
  [[nodiscard]] TENZR_API Array Sqrt(const Array& x);
  [[nodiscard]] TENZR_API Array Rsqrt(const Array& x);
  [[nodiscard]] TENZR_API Array Sin(const Array& x);
  [[nodiscard]] TENZR_API Array Cos(const Array& x);
  [[nodiscard]] TENZR_API Array Tan(const Array& x);
  [[nodiscard]] TENZR_API Array Ceil(const Array& x);
  [[nodiscard]] TENZR_API Array Floor(const Array& x);
  [[nodiscard]] TENZR_API Array Round(const Array& x);
  [[nodiscard]] TENZR_API Array Sign(const Array& x);

  // Reductions

  [[nodiscard]] TENZR_API Array ReduceSum(const Array& operand, const std::vector<int64_t>& axes = {}, bool keepdims = false);

  [[nodiscard]] TENZR_API Array ReduceMax(const Array& operand, const std::vector<int64_t>& axes = {}, bool keepdims = false);

  [[nodiscard]] TENZR_API Array ReduceMin(const Array& operand, const std::vector<int64_t>& axes = {}, bool keepdims = false);

  [[nodiscard]] TENZR_API Array ReduceProd(const Array& operand, const std::vector<int64_t>& axes = {}, bool keepdims = false);

  [[nodiscard]] TENZR_API Array ReduceWindow(const Array& operand, const Array& init_value, const std::vector<int64_t>& window_dimensions,
                                             const std::vector<int64_t>& window_strides = {}, const std::vector<int64_t>& padding_low = {},
                                             const std::vector<int64_t>& padding_high = {}, std::string_view reducer = "add");

  // Windowing, Slicing & Padding

  [[nodiscard]] TENZR_API Array slice(const Array& operand, const std::vector<int64_t>& start_indices, const std::vector<int64_t>& limit_indices,
                                      const std::vector<int64_t>& strides = {});

  [[nodiscard]] TENZR_API Array Pad(const Array& operand, const Array& padding_value, const std::vector<int64_t>& edge_padding_low,
                                    const std::vector<int64_t>& edge_padding_high, const std::vector<int64_t>& interior_padding = {});

  // Indexing

  [[nodiscard]] TENZR_API Array Gather(const Array& operand, const Array& start_indices, int64_t axis = 0);

  [[nodiscard]] TENZR_API Array DynamicSlice(const Array& operand, const Array& start_indices, const std::vector<int64_t>& slice_sizes);

  [[nodiscard]] TENZR_API Array DynamicSlice(const Array& operand, const std::vector<Array>& start_indices, const Shape& slice_sizes);

  [[nodiscard]] TENZR_API Array DynamicUpdateSlice(const Array& operand, const Array& update, const Array& start_indices);

  [[nodiscard]] TENZR_API Array DynamicUpdateSlice(const Array& operand, const Array& update, const std::vector<Array>& start_indices);

  [[nodiscard]] TENZR_API Array ScatterAdd(const Array& operand, const Array& indices, const Array& updates, int64_t axis = 0);

  [[nodiscard]] TENZR_API Array Scatter(const Array& operand, const Array& indices, const Array& updates, int64_t axis = 0, std::string_view update = "add");

  [[nodiscard]] TENZR_API Array Argsort(const Array& operand, int64_t axis = -1, bool descending = false);

  [[nodiscard]] TENZR_API Array Sort(const Array& operand, int64_t axis = -1, bool descending = false);

  [[nodiscard]] TENZR_API Array Argmax(const Array& operand, int64_t axis = -1);

  [[nodiscard]] TENZR_API Array Argmin(const Array& operand, int64_t axis = -1);

  [[nodiscard]] TENZR_API Array Concatenate(const std::vector<Array>& arrays, int64_t axis = 0);

  [[nodiscard]] TENZR_API Array Stack(const std::vector<Array>& arrays, int64_t axis = 0);

  [[nodiscard]] TENZR_API Array Conv2d(const Array& input, const Array& weight, std::pair<int64_t, int64_t> stride = {1, 1},
                                       std::pair<int64_t, int64_t> padding = {0, 0});

  [[nodiscard]] TENZR_API Array ConvGeneral(const Array& lhs, const Array& rhs, const std::vector<int64_t>& window_strides,
                                            const std::vector<int64_t>& lhs_dilation, const std::vector<int64_t>& rhs_dilation,
                                            const std::vector<int64_t>& padding_low, const std::vector<int64_t>& padding_high,
                                            const std::vector<int64_t>& lhs_spec = {0, 1, 2, 3}, const std::vector<int64_t>& rhs_spec = {0, 1, 2, 3},
                                            const std::vector<int64_t>& out_spec = {0, 1, 2, 3});

  [[nodiscard]] TENZR_API Array TakeAlongAxis(const Array& operand, const Array& indices, int64_t axis = -1);

  [[nodiscard]] TENZR_API Array PutAlongAxis(const Array& operand, const Array& indices, const Array& values, int64_t axis = -1,
                                             std::string_view update = "set");

  [[nodiscard]] TENZR_API TopKResult TopK(const Array& operand, int64_t k, int64_t axis = -1, bool largest = true);

  [[nodiscard]] TENZR_API Array Reverse(const Array& operand, const std::vector<int64_t>& dimensions);

  // Tensor Contraction & Structural Primitives

  [[nodiscard]] TENZR_API Array DotGeneral(const Array& lhs, const Array& rhs, const std::vector<int64_t>& lhs_contracting_dims,
                                           const std::vector<int64_t>& rhs_contracting_dims, const std::vector<int64_t>& lhs_batch_dims = {},
                                           const std::vector<int64_t>& rhs_batch_dims = {});

  [[nodiscard]] TENZR_API Array BroadcastInDim(const Array& operand, const Shape& target_shape, const std::vector<int64_t>& broadcast_dimensions);

  [[nodiscard]] TENZR_API Array Reshape(const Array& operand, const Shape& new_shape);

  [[nodiscard]] TENZR_API Array Transpose(const Array& operand, const std::vector<int64_t>& permutation);

  [[nodiscard]] TENZR_API Array ConvertElementType(const Array& operand, DType new_dtype);

  [[nodiscard]] TENZR_API Array Select(const Array& pred, const Array& on_true, const Array& on_false);

  [[nodiscard]] TENZR_API Array Clamp(const Array& min_val, const Array& operand, const Array& max_val);

  // Distributed Collectives & Sharding Primitives

  [[nodiscard]] TENZR_API Array AllReduce(const Array& x, std::string_view op = "sum", const std::vector<std::vector<int64_t>>& replica_groups = {});

  [[nodiscard]] TENZR_API Array AllGather(const Array& x, int64_t all_gather_dim = 0, const std::vector<std::vector<int64_t>>& replica_groups = {});

  [[nodiscard]] TENZR_API Array AllToAll(const Array& x, int64_t split_dim = 0, int64_t concat_dim = 0, int64_t split_count = 1,
                                         const std::vector<std::vector<int64_t>>& replica_groups = {});

  [[nodiscard]] TENZR_API Array ReduceScatter(const Array& x, std::string_view op = "sum", int64_t scatter_dim = 0,
                                              const std::vector<std::vector<int64_t>>& replica_groups = {});

  [[nodiscard]] TENZR_API Array Psum(const Array& x, std::string_view axis_name = "");

  [[nodiscard]] TENZR_API Array Pmean(const Array& x, std::string_view axis_name = "");

  [[nodiscard]] TENZR_API Array WithShardingConstraint(const Array& x, const NamedSharding& sharding);

}  // namespace tenzr::plax

#endif  // TENZR_PLAX_PLAX_H_