src/txpr/control_flow.cc¶

// 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.

#include "tenzr/txpr/control_flow.h"

#include <string>
#include <utility>
#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/numpy/numpy.h"
#include "tenzr/plax/plax.h"
#include "tenzr/txpr/tracer.h"

namespace tenzr::txpr {

  Array Cond(const Array& pred, const std::function<Array()>& true_fn, const std::function<Array()>& false_fn) {
    TENZR_CHECK(pred.Dtype() == DType::kBool, "cond predicate must have DType::Bool");

    if (pred.HasData() && !TracerContext::IsActive()) {
      bool p = pred.Item<bool>();
      return p ? true_fn() : false_fn();
    }

    // Under tracing, trace both branches and emit selection
    Array res_true = true_fn();
    Array res_false = false_fn();
    TENZR_CHECK(res_true.shape() == res_false.shape(), "cond branches must return arrays with identical shapes");
    TENZR_CHECK(res_true.Dtype() == res_false.Dtype(), "cond branches must return arrays with identical dtypes");

    Array broadcast_pred = pred.BroadcastTo(res_true.shape());
    return plax::Select(broadcast_pred, res_true, res_false);
  }

  Array Cond(const Array& pred, const std::function<Array(const Array&)>& true_fn, const std::function<Array(const Array&)>& false_fn, const Array& operand) {
    return Cond(pred, [&]() { return true_fn(operand); }, [&]() { return false_fn(operand); });
  }

  ScanResult Scan(const std::function<ScanResult(const Array& carry, const Array& x)>& body_fn, const Array& init_carry, const Array& xs, int64_t length) {
    TENZR_CHECK(xs.shape().Rank() >= 1, "scan xs must have rank at least 1");
    int64_t n = (length >= 0) ? length : xs.shape().Dim(0);
    TENZR_CHECK(n >= 0, "scan length must be non-negative");

    Array cur_carry = init_carry;
    std::vector<Array> ys;
    ys.reserve(static_cast<size_t>(n));

    for (int64_t i = 0; i < n; ++i) {
      Array x_i = xs.slice({Slice(i, i + 1, 1)}).Squeeze(0).Contiguous();
      auto [next_carry, y_i] = body_fn(cur_carry, x_i);
      cur_carry = next_carry;
      ys.push_back(y_i);
    }

    Array stacked_ys = tnp::Stack(ys, 0);
    return {cur_carry, stacked_ys};
  }

  Array WhileLoop(const std::function<Array(const Array&)>& cond_fn, const std::function<Array(const Array&)>& body_fn, const Array& init_val) {
    Array cur_val = init_val;
    int64_t iteration_count = 0;
    const int64_t kMaxWhileIterations = 100000;

    while (true) {
      Array cond_res = cond_fn(cur_val);
      TENZR_CHECK(cond_res.Dtype() == DType::kBool, "while_loop cond_fn must return DType::Bool");
      TENZR_CHECK(cond_res.shape().NumElements() == 1, "while_loop cond_fn must return scalar boolean");

      if (!cond_res.Item<bool>()) {
        break;
      }

      cur_val = body_fn(cur_val);
      iteration_count++;
      TENZR_CHECK(iteration_count < kMaxWhileIterations, "while_loop exceeded maximum iteration limit (100,000)");
    }

    return cur_val;
  }

}  // namespace tenzr::txpr