// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.

//! # rlx-tensor
//!
//! Symbolic tensor API for RLX — expression syntax with compiler-backed
//! execution. Compared to NumPy / ndarray / PyTorch / JAX:
//!
//! | | ndarray / NumPy | PyTorch | JAX | **rlx-tensor** |
//! |-|-----------------|---------|-----|----------------|
//! | Default execution | eager, host | eager, GPU | trace + XLA | **trace → fuse → backend** |
//! | Intermediate buffers | always materialized | often kept for autograd | XLA decides | **fusion + memory plan** |
//! | Device | host only | CUDA-centric | TPU/GPU | **CPU / Metal / CUDA / ROCm / wgpu / MLX / TPU** |
//! | Variable batch/seq | manual padding | manual | `jax.vmap` / shapes | **first-class `Dim::Dynamic`** |
//! | Slicing | views (cheap) | views | slice ops | **`s![]` → narrow (fusable)** |
//!
//! Symbolic [`Tensor`] handles carry **zero payload bytes** until you compile
//! and run — but with the `eval` feature you can also build from host data
//! NumPy-style and read results straight back, the graph staying lazy (so the
//! compiler fuses + memory-plans the whole expression) until you materialize:
//!
//! ```ignore
//! use rlx_tensor::Tensor; // requires feature = "eval"
//!
//! let a = Tensor::from_vec(vec![1.0, 2.0, 3.0], [3]);
//! let b = Tensor::ones([3]);
//! let c = (&a + &b).relu();        // graphs auto-merge, nothing runs yet
//! assert_eq!(c.to_vec(), vec![2.0, 3.0, 4.0]); // auto-selects fastest backend
//! // c.on(Device::Metal).to_vec()  // or pin an explicit backend
//! ```
//!
//! Reverse-mode autodiff (feature `grad`) returns gradient tensors:
//!
//! ```ignore
//! let loss = (&a * &b).sum([0], false);
//! let g = loss.grad(&[&a, &b]);     // [∂loss/∂a, ∂loss/∂b]
//! assert_eq!(g[0].to_vec(), b.to_vec());
//! ```
//!
//! JAX-shaped composable transforms (feature `transforms`) operate on [`Func`]
//! (a traced function of named inputs) and chain — `vmap(grad(f))`:
//!
//! ```ignore
//! let f = Func::new("f", |s| { let x = s.input("x", shape![3]); (&x * &x).sum([0], false) });
//! let batched_grad = f.grad(&["x"]).vmap(&["x"], 4);  // Func -> Func -> Func
//! let compiled = batched_grad.jit();                  // compile once
//! let out = compiled.run(&[("x", &xs)]);              // run many, no recompile
//! ```
//!
//! Materialization is compile-cached per thread (keyed by graph + constants +
//! device), so even plain `to_vec` / `run` reuse the compiled artifact.
//!
//! ```rust
//! use rlx_tensor::{ax, graph, rg, s, shape};
//!
//! // static shape
//! let g = graph("mlp", |g| {
//!     let x = g.input("x", shape![2, 4]);
//!     let w = g.param("w", shape![4, 3]);
//!     (&x.matmul(&w)).gelu()
//! });
//!
//! // dynamic batch + window slice (no copy in IR)
//! let g2 = graph("batched", |g| {
//!     let x = g.input("x", shape![?, 128]);
//!     x.slice(s![ax(), rg(0, 64)])
//! });
//! # let _ = (g, g2);
//! ```

mod array;
#[cfg(feature = "eval")]
mod cache;
mod handle;
#[cfg(feature = "ndarray")]
mod interop;
mod scalar;
#[cfg(feature = "optim")]
mod schedule;
mod scope;
mod slice;
mod tensor;
mod transform;

pub use array::{cat, stack};
pub use rlx_ir::op::MaskKind;
pub use rlx_ir::{DType, Dim, Graph, NodeId, Shape};
/// Optimizers for [`Func::train_step`] (re-exported from `rlx_optim`).
/// Available with the `optim` feature.
#[cfg(feature = "optim")]
pub use rlx_optim::{Adam, AdamW, Lion, Optimizer, Sgd};
#[cfg(feature = "optim")]
pub use schedule::LrSchedule;
pub use scope::{GraphScope, graph, graph_with};
pub use slice::{SliceAxis, SliceSpec, ax, ix, rg, tail};
pub use tensor::Tensor;
pub use transform::Func;
#[cfg(feature = "eval")]
pub use transform::Jitted;

#[cfg(feature = "eval")]
pub use cache::{cache_stats, clear_cache};
/// Device selector for [`Tensor::on`] / [`Tensor::to_vec_on`] (re-exported
/// from `rlx_runtime`). Available with the `eval` feature.
#[cfg(feature = "eval")]
pub use rlx_runtime::Device;
/// Backend detection (re-exported from `rlx_runtime`): which devices this build
/// can use, and the fastest one. `to_vec` / `Func::run` pick automatically.
#[cfg(feature = "eval")]
pub use rlx_runtime::{available_devices, fastest_device, fastest_device_for, is_available};
#[cfg(feature = "eval")]
pub use tensor::Materialize;

#[doc(hidden)]
#[macro_export]
macro_rules! __dim {
    (?) => {
        $crate::Dim::Dynamic(0)
    };
    (? $sym:literal) => {
        $crate::Dim::Dynamic($sym)
    };
    ($n:expr) => {
        $crate::Dim::Static($n)
    };
}

/// Build a [`Shape`]. Default dtype is `F32`.
///
/// Dynamic dimensions: `?` → symbol 0 (repeat `?` for the same unknown size),
/// `?1` → symbol 1, etc.
///
/// ```rust
/// use rlx_tensor::{shape, DType};
///
/// let a = shape![2, 4];
/// let b = shape![F32; ?, 128];
/// assert_eq!(b.dtype(), DType::F32);
/// assert!(!b.is_static());
/// ```
#[macro_export]
macro_rules! shape {
    ($dtype:ident; $($d:tt),+ $(,)?) => {
        $crate::Shape::from_dims(
            &[$( $crate::__dim!($d) ),+],
            $crate::DType::$dtype,
        )
    };
    ($($d:tt),+ $(,)?) => {
        $crate::Shape::from_dims(
            &[$( $crate::__dim!($d) ),+],
            $crate::DType::F32,
        )
    };
}
