tenferro_ad/
eager_ops.rs

1use std::sync::Arc;
2
3use computegraph::GraphOperation;
4use num_complex::{Complex32, Complex64};
5use tenferro_ops::broadcast::{
6    broadcast_error_to_validation, broadcast_in_dim_extent_error, broadcast_input_plan,
7    broadcast_shape, broadcast_shapes,
8};
9use tenferro_ops::dim_expr::DimExpr;
10use tenferro_ops::std_tensor_op::StdTensorOp;
11use tenferro_tensor::{
12    DType, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig, Tensor,
13    TensorValue,
14};
15
16use crate::eager::{
17    eager_capture_active, eager_grad_recording_enabled, eager_op_profile_start, exec_single_output,
18    exec_single_output_read, maybe_print_eager_op_profile, profile_eager_op_section,
19    record_eager_op_profile, record_eager_outputs, record_eager_value_outputs, EagerTensor,
20};
21use crate::eager_exec::exec_dot_general_with_conj_on_tensor_reads;
22use crate::error::{Error, Result};
23
24pub(crate) fn broadcast_binary(
25    op: &'static str,
26    lhs: &EagerTensor,
27    rhs: &EagerTensor,
28) -> Result<(EagerTensor, EagerTensor)> {
29    ensure_same_context(lhs, rhs)?;
30    let shape =
31        broadcast_shape(lhs.shape(), rhs.shape()).map_err(|err| broadcast_error(op, err))?;
32    Ok((
33        broadcast_to(op, lhs, &shape)?,
34        broadcast_to(op, rhs, &shape)?,
35    ))
36}
37
38pub(crate) fn broadcast_ternary(
39    op: &'static str,
40    first: &EagerTensor,
41    second: &EagerTensor,
42    third: &EagerTensor,
43) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
44    ensure_same_context(first, second)?;
45    ensure_same_context(first, third)?;
46    let shape = broadcast_shapes([first.shape(), second.shape(), third.shape()])
47        .map_err(|err| broadcast_error(op, err))?;
48    Ok((
49        broadcast_to(op, first, &shape)?,
50        broadcast_to(op, second, &shape)?,
51        broadcast_to(op, third, &shape)?,
52    ))
53}
54
55fn broadcast_to(
56    op: &'static str,
57    input: &EagerTensor,
58    target_shape: &[usize],
59) -> Result<EagerTensor> {
60    let input_shape = input.shape();
61    if input_shape == target_shape {
62        return Ok(input.clone());
63    }
64
65    let plan =
66        broadcast_input_plan(input_shape, target_shape).map_err(|err| broadcast_error(op, err))?;
67    let source = if plan.source_shape == input_shape {
68        input.clone()
69    } else {
70        input.reshape(&plan.source_shape)?
71    };
72    source.broadcast_in_dim(target_shape, &plan.dims)
73}
74
75fn broadcast_error(op: &'static str, err: tenferro_ops::broadcast::BroadcastError) -> Error {
76    tenferro_tensor::Error::validation(op, broadcast_error_to_validation(err)).into()
77}
78
79fn ensure_same_context(lhs: &EagerTensor, rhs: &EagerTensor) -> Result<()> {
80    if !lhs.same_context(rhs) {
81        return Err(Error::ContextMismatch {
82            lhs: lhs.ctx_id(),
83            rhs: rhs.ctx_id(),
84        });
85    }
86    Ok(())
87}
88
89impl std::ops::Add for &EagerTensor {
90    type Output = Result<EagerTensor>;
91
92    fn add(self, rhs: &EagerTensor) -> Result<EagerTensor> {
93        EagerTensor::add(self, rhs)
94    }
95}
96
97impl std::ops::Sub for &EagerTensor {
98    type Output = Result<EagerTensor>;
99
100    fn sub(self, rhs: &EagerTensor) -> Result<EagerTensor> {
101        EagerTensor::sub(self, rhs)
102    }
103}
104
105impl std::ops::Mul for &EagerTensor {
106    type Output = Result<EagerTensor>;
107
108    fn mul(self, rhs: &EagerTensor) -> Result<EagerTensor> {
109        EagerTensor::mul(self, rhs)
110    }
111}
112
113impl std::ops::Div for &EagerTensor {
114    type Output = Result<EagerTensor>;
115
116    fn div(self, rhs: &EagerTensor) -> Result<EagerTensor> {
117        EagerTensor::div(self, rhs)
118    }
119}
120
121impl std::ops::Rem for &EagerTensor {
122    type Output = Result<EagerTensor>;
123
124    fn rem(self, rhs: &EagerTensor) -> Result<EagerTensor> {
125        EagerTensor::rem(self, rhs)
126    }
127}
128
129impl std::ops::Neg for &EagerTensor {
130    type Output = Result<EagerTensor>;
131
132    fn neg(self) -> Result<EagerTensor> {
133        EagerTensor::neg(self)
134    }
135}
136
137impl EagerTensor {
138    /// Elementwise addition.
139    ///
140    /// # Examples
141    ///
142    /// ```
143    /// use tenferro_cpu::CpuBackend;
144    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
145    ///
146    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
147    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
148    /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone()).unwrap();
149    /// let z = x.add(&y).unwrap();
150    ///
151    /// assert_eq!(z.value().unwrap().as_slice::<f64>().unwrap(), &[4.0, 6.0]);
152    /// # Ok::<(), tenferro_ad::Error>(())
153    /// ```
154    ///
155    /// # Errors
156    ///
157    /// Returns [`Error::ContextMismatch`] for tensors from different eager
158    /// runtimes, [`tenferro_tensor::Error::Validation`] with
159    /// `ShapeMismatch`/`DTypeMismatch` for incompatible operands, or a typed
160    /// backend/runtime-state error during execution.
161    pub fn add(&self, other: &Self) -> Result<Self> {
162        let (lhs, rhs) = broadcast_binary("add", self, other)?;
163        lhs.binary_op(&rhs, StdTensorOp::Add)
164    }
165
166    /// Elementwise subtraction.
167    ///
168    /// # Errors
169    ///
170    /// Returns [`Error::ContextMismatch`] for tensors from different eager
171    /// runtimes, [`tenferro_tensor::Error::Validation`] with
172    /// `ShapeMismatch`/`DTypeMismatch` for incompatible operands, or a typed
173    /// backend/runtime-state error during execution.
174    pub fn sub(&self, other: &Self) -> Result<Self> {
175        let (lhs, rhs) = broadcast_binary("sub", self, other)?;
176        lhs.binary_op(&rhs, StdTensorOp::Sub)
177    }
178
179    /// Elementwise multiplication.
180    ///
181    /// # Examples
182    ///
183    /// ```
184    /// use tenferro_cpu::CpuBackend;
185    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
186    ///
187    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
188    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
189    /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone()).unwrap();
190    /// let z = x.mul(&y).unwrap();
191    ///
192    /// assert_eq!(z.value().unwrap().as_slice::<f64>().unwrap(), &[3.0, 8.0]);
193    /// # Ok::<(), tenferro_ad::Error>(())
194    /// ```
195    ///
196    /// # Errors
197    ///
198    /// Returns [`Error::ContextMismatch`] for tensors from different eager
199    /// runtimes, [`tenferro_tensor::Error::Validation`] with
200    /// `ShapeMismatch`/`DTypeMismatch` for incompatible operands, or a typed
201    /// backend/runtime-state error during execution.
202    pub fn mul(&self, other: &Self) -> Result<Self> {
203        let (lhs, rhs) = broadcast_binary("mul", self, other)?;
204        lhs.binary_op(&rhs, StdTensorOp::Mul)
205    }
206
207    /// Negate the tensor.
208    ///
209    /// # Examples
210    ///
211    /// ```
212    /// use tenferro_cpu::CpuBackend;
213    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
214    ///
215    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
216    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, -2.0]).unwrap(), ctx.clone()).unwrap();
217    /// let y = x.neg().unwrap();
218    ///
219    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[-1.0, 2.0]);
220    /// # Ok::<(), tenferro_ad::Error>(())
221    /// ```
222    ///
223    /// # Errors
224    ///
225    /// Returns [`tenferro_tensor::Error::Unsupported`] when the backend does
226    /// not implement negation for the dtype, or a typed backend/runtime-state
227    /// error during execution.
228    pub fn neg(&self) -> Result<Self> {
229        self.unary_op(StdTensorOp::Neg)
230    }
231
232    /// Elementwise exponential.
233    ///
234    /// # Examples
235    ///
236    /// ```
237    /// use tenferro_cpu::CpuBackend;
238    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
239    ///
240    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
241    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![0.0_f64]).unwrap(), ctx.clone()).unwrap();
242    /// let y = x.exp().unwrap();
243    ///
244    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[1.0]);
245    /// # Ok::<(), tenferro_ad::Error>(())
246    /// ```
247    ///
248    /// # Errors
249    ///
250    /// Returns [`tenferro_tensor::Error::Unsupported`] when the backend does
251    /// not implement exponentiation for the dtype, or a typed backend/
252    /// runtime-state error during execution.
253    pub fn exp(&self) -> Result<Self> {
254        self.unary_op(StdTensorOp::Exp)
255    }
256
257    /// Reduce sum over the requested axes.
258    ///
259    /// # Examples
260    ///
261    /// ```
262    /// use tenferro_cpu::CpuBackend;
263    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
264    ///
265    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
266    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap(), ctx.clone()).unwrap();
267    /// let y = x.reduce_sum(None).unwrap();
268    ///
269    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[10.0]);
270    /// # Ok::<(), tenferro_ad::Error>(())
271    /// ```
272    ///
273    /// # Errors
274    ///
275    /// Returns [`tenferro_tensor::Error::Validation`] with `AxisOutOfBounds` or
276    /// `DuplicateAxis` for an invalid reduction axis, or a typed
277    /// unsupported/backend/runtime-state error for the selected dtype.
278    pub fn reduce_sum(&self, axes: Option<&[usize]>) -> Result<Self> {
279        let axes = axes.map_or_else(|| (0..self.shape().len()).collect(), <[usize]>::to_vec);
280        validate_eager_axes("EagerTensor::reduce_sum", self.shape().len(), &axes)?;
281        self.unary_op(StdTensorOp::ReduceSum { axes })
282    }
283
284    /// Sum elementwise squares over the requested axes.
285    ///
286    /// Each value is squared in its input dtype before reduction. The initial
287    /// supported dtypes are `f32` and `f64`; other dtypes return a typed
288    /// unsupported error. Passing an empty axis slice returns the elementwise
289    /// square without reducing rank.
290    ///
291    /// This operation is useful when the squared sum is needed directly. Use
292    /// the linalg norm APIs when a square root or complex magnitude semantics
293    /// are required.
294    ///
295    /// # Errors
296    ///
297    /// Returns a typed validation error for invalid axes, a typed unsupported
298    /// error for other dtypes, or a typed backend or runtime-state error during
299    /// execution.
300    pub fn reduce_sum_squares(&self, axes: &[usize]) -> Result<Self> {
301        validate_eager_axes("EagerTensor::reduce_sum_squares", self.shape().len(), axes)?;
302        self.unary_op(StdTensorOp::ReduceSumSquares {
303            axes: axes.to_vec(),
304        })
305    }
306
307    /// Execute a dot-general contraction eagerly.
308    ///
309    /// # Examples
310    ///
311    /// ```
312    /// use tenferro_cpu::CpuBackend;
313    /// use tenferro_ad::{DotGeneralConfig, EagerRuntime, EagerTensor, Tensor};
314    ///
315    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
316    /// let a = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap(), ctx.clone()).unwrap();
317    /// let b = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![3, 2], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap(), ctx.clone()).unwrap();
318    /// let c = a.dot_general(&b, DotGeneralConfig {
319    ///     lhs_contracting_dims: vec![1],
320    ///     rhs_contracting_dims: vec![0],
321    ///     lhs_batch_dims: vec![],
322    ///     rhs_batch_dims: vec![],
323    /// }).unwrap();
324    ///
325    /// assert_eq!(c.shape(), &[2, 2]);
326    /// # Ok::<(), tenferro_ad::Error>(())
327    /// ```
328    ///
329    /// # Errors
330    ///
331    /// Returns [`tenferro_tensor::Error::Validation`] with `RankMismatch`,
332    /// `AxisOutOfBounds`, `DuplicateAxis`, `ShapeMismatch`, or `DTypeMismatch`
333    /// when `config` or the operands are invalid; backend and runtime-state
334    /// failures retain their typed sources.
335    pub fn dot_general(&self, other: &Self, config: DotGeneralConfig) -> Result<Self> {
336        validate_eager_dot_general_config(
337            "EagerTensor::dot_general",
338            &config,
339            self.shape().len(),
340            other.shape().len(),
341        )?;
342        self.binary_op(other, StdTensorOp::DotGeneral { config })
343    }
344
345    /// Execute a dot-general contraction, optionally conjugating either operand.
346    ///
347    /// Untracked tensors route the conjugation flags directly to the backend so
348    /// the conjugated operand does not need to be materialized. Tracked tensors
349    /// fall back to explicit `Conj` plus `DotGeneral` so reverse-mode AD keeps
350    /// the same graph semantics as the standard eager ops.
351    ///
352    /// # Errors
353    ///
354    /// Returns [`Error::ContextMismatch`] for operands from different eager
355    /// runtimes, [`tenferro_tensor::Error::Validation`] for rank/axis/shape or
356    /// dtype mismatches in `config`, or a typed backend/runtime-state error.
357    pub fn dot_general_with_conj(
358        &self,
359        other: &Self,
360        config: DotGeneralConfig,
361        lhs_conj: bool,
362        rhs_conj: bool,
363    ) -> Result<Self> {
364        if !self.same_context(other) {
365            return Err(Error::ContextMismatch {
366                lhs: self.ctx_id(),
367                rhs: other.ctx_id(),
368            });
369        }
370        validate_eager_dot_general_config(
371            "EagerTensor::dot_general_with_conj",
372            &config,
373            self.shape().len(),
374            other.shape().len(),
375        )?;
376
377        if !self.requires_grad && !other.requires_grad {
378            let ctx = Arc::clone(&self.ctx);
379            let mut backend = ctx.lock_backend()?;
380            let output = exec_dot_general_with_conj_on_tensor_reads(
381                self.tensor_read(),
382                other.tensor_read(),
383                &config,
384                lhs_conj,
385                rhs_conj,
386                &mut *backend,
387            )?;
388            drop(backend);
389            return Self::new_untracked_result(ctx, output);
390        }
391
392        match (lhs_conj, rhs_conj) {
393            (false, false) => self.dot_general(other, config),
394            (true, false) => self.conj()?.dot_general(other, config),
395            (false, true) => {
396                let rhs = other.conj()?;
397                self.dot_general(&rhs, config)
398            }
399            (true, true) => {
400                let lhs = self.conj()?;
401                let rhs = other.conj()?;
402                lhs.dot_general(&rhs, config)
403            }
404        }
405    }
406
407    /// Scale by a real scalar: `y = factor * x`.
408    ///
409    /// Integer factors are rounded to the nearest integer before multiplication,
410    /// boolean factors map finite zero to `false` and other finite values to
411    /// `true`, and complex tensors receive a zero-imaginary scalar.
412    ///
413    /// # Errors
414    ///
415    /// Returns [`Error::TensorRuntime`] with
416    /// [`tenferro_tensor::ValidationError::InvalidArgument`] when an integer or
417    /// boolean factor is non-finite or outside the input dtype's range. Backend
418    /// and runtime execution failures retain their typed source variants.
419    /// # Examples
420    ///
421    /// ```rust
422    /// # use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
423    /// # use tenferro_cpu::CpuBackend;
424    /// # let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
425    /// # let x = EagerTensor::from_tensor_in(
426    /// #     Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(),
427    /// #     ctx,
428    /// # )?;
429    /// let scaled = x.scale_real(2.0)?;
430    /// assert_eq!(scaled.value()?.as_slice::<f64>()?, &[2.0, 4.0]);
431    /// # Ok::<(), tenferro_ad::Error>(())
432    /// ```
433    pub fn scale_real(&self, factor: f64) -> Result<Self> {
434        let scalar = match self.dtype() {
435            DType::F64 => Tensor::from_vec_col_major(vec![], vec![factor])?,
436            DType::F32 => Tensor::from_vec_col_major(vec![], vec![factor as f32])?,
437            DType::I32 => Tensor::from_vec_col_major(vec![], vec![round_real_to_i32(factor)?])?,
438            DType::I64 => Tensor::from_vec_col_major(vec![], vec![round_real_to_i64(factor)?])?,
439            DType::Bool => Tensor::from_vec_col_major(vec![], vec![bool_from_real(factor)?])?,
440            DType::C64 => Tensor::from_vec_col_major(vec![], vec![Complex64::new(factor, 0.0)])?,
441            DType::C32 => {
442                Tensor::from_vec_col_major(vec![], vec![Complex32::new(factor as f32, 0.0)])?
443            }
444        };
445        let scalar = EagerTensor::from_tensor_in(scalar, Arc::clone(&self.ctx))?;
446        self.mul(&scalar)
447    }
448
449    /// Scale a complex tensor by a complex scalar: `y = factor * x`.
450    ///
451    /// # Errors
452    ///
453    /// Returns [`Error::TensorRuntime`] with
454    /// [`tenferro_tensor::ValidationError::InvalidArgument`] for a non-complex
455    /// input dtype. Backend and runtime execution failures retain their typed
456    /// source variants.
457    /// # Examples
458    ///
459    /// ```rust
460    /// # use num_complex::Complex64;
461    /// # use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
462    /// # use tenferro_cpu::CpuBackend;
463    /// # let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
464    /// # let x = EagerTensor::from_tensor_in(
465    /// #     Tensor::from_vec_col_major(
466    /// #         vec![2],
467    /// #         vec![Complex64::new(1.0, 2.0), Complex64::new(3.0, 4.0)],
468    /// #     )
469    /// #     .unwrap(),
470    /// #     ctx,
471    /// # )?;
472    /// let scaled = x.scale_complex(Complex64::new(0.0, 1.0))?;
473    /// assert_eq!(
474    ///     scaled.value()?.as_slice::<Complex64>()?,
475    ///     &[Complex64::new(-2.0, 1.0), Complex64::new(-4.0, 3.0)],
476    /// );
477    /// # Ok::<(), tenferro_ad::Error>(())
478    /// ```
479    pub fn scale_complex(&self, factor: Complex64) -> Result<Self> {
480        let scalar = match self.dtype() {
481            DType::C64 => Tensor::from_vec_col_major(vec![], vec![factor])?,
482            DType::C32 => Tensor::from_vec_col_major(
483                vec![],
484                vec![Complex32::new(factor.re as f32, factor.im as f32)],
485            )?,
486            dtype => {
487                return Err(Error::TensorRuntime(
488                    tenferro_tensor::Error::invalid_argument(
489                        "scale_complex",
490                        "dtype",
491                        format!("requires complex tensor dtype, got {dtype:?}"),
492                    ),
493                ));
494            }
495        };
496        let scalar = EagerTensor::from_tensor_in(scalar, Arc::clone(&self.ctx))?;
497        self.mul(&scalar)
498    }
499
500    /// Matrix multiplication for rank-2 tensors.
501    ///
502    /// This is a convenience wrapper over [`Self::dot_general`] that
503    /// contracts the left matrix's column axis with the right matrix's row
504    /// axis.
505    ///
506    /// # Examples
507    ///
508    /// ```
509    /// use tenferro_cpu::CpuBackend;
510    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
511    ///
512    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
513    /// let a = EagerTensor::from_tensor_in(
514    ///     Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap(),
515    ///     ctx.clone(),
516    /// ).unwrap();
517    /// let b = EagerTensor::from_tensor_in(
518    ///     Tensor::from_vec_col_major(vec![2, 1], vec![5.0_f64, 6.0]).unwrap(),
519    ///     ctx,
520    /// ).unwrap();
521    /// let c = a.matmul(&b).unwrap();
522    ///
523    /// assert_eq!(c.shape(), &[2, 1]);
524    /// assert_eq!(c.value().unwrap().as_slice::<f64>().unwrap(), &[23.0, 34.0]);
525    /// # Ok::<(), tenferro_ad::Error>(())
526    /// ```
527    ///
528    /// # Errors
529    ///
530    /// Returns [`tenferro_tensor::ValidationError::RankMismatch`] when either operand is
531    /// not rank 2, `ShapeMismatch` when the inner dimensions differ, or a typed
532    /// dtype/backend/runtime-state error during the contraction.
533    pub fn matmul(&self, other: &Self) -> Result<Self> {
534        let lhs_shape = self.shape();
535        let rhs_shape = other.shape();
536        if lhs_shape.len() != 2 {
537            return Err(tenferro_tensor::Error::rank_mismatch("matmul", 2, lhs_shape.len()).into());
538        }
539        if rhs_shape.len() != 2 {
540            return Err(tenferro_tensor::Error::rank_mismatch("matmul", 2, rhs_shape.len()).into());
541        }
542        if lhs_shape[1] != rhs_shape[0] {
543            return Err(
544                tenferro_tensor::Error::shape_mismatch("matmul", lhs_shape, rhs_shape).into(),
545            );
546        }
547        self.dot_general(
548            other,
549            DotGeneralConfig {
550                lhs_contracting_dims: vec![1],
551                rhs_contracting_dims: vec![0],
552                lhs_batch_dims: vec![],
553                rhs_batch_dims: vec![],
554            },
555        )
556    }
557
558    /// Permute tensor axes.
559    ///
560    /// # Examples
561    ///
562    /// ```
563    /// use tenferro_cpu::CpuBackend;
564    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
565    ///
566    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
567    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(
568    ///     vec![2, 3],
569    ///     vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0],
570    /// ).unwrap(), ctx.clone()).unwrap();
571    /// let y = x.transpose(&[1, 0]).unwrap();
572    ///
573    /// assert_eq!(y.shape(), &[3, 2]);
574    /// let materialized = y.to_tensor().unwrap();
575    /// assert_eq!(materialized.as_slice::<f64>().unwrap(), &[1.0, 3.0, 5.0, 2.0, 4.0, 6.0]);
576    /// # Ok::<(), tenferro_ad::Error>(())
577    /// ```
578    ///
579    /// # Errors
580    ///
581    /// Returns [`tenferro_tensor::Error::Validation`] with `AxisOutOfBounds`
582    /// or `DuplicateAxis` when `perm` is not a permutation, or a typed
583    /// backend/runtime-state error while creating the view.
584    pub fn transpose(&self, perm: &[usize]) -> Result<Self> {
585        let op = StdTensorOp::Transpose {
586            perm: perm.to_vec(),
587        };
588        // INVARIANT: the result must own a group independent of `self`; the
589        // explicit duplicate is the ownership boundary before making a view.
590        let base = self.to_tensor()?;
591        let value = TensorValue::from_tensor(base)
592            .transpose_view(perm)
593            .map_err(Error::TensorRuntime)?;
594        Self::nary_value_op(&[self], op, value)
595    }
596
597    /// Reshape without changing element order.
598    ///
599    /// # Examples
600    ///
601    /// ```
602    /// use tenferro_cpu::CpuBackend;
603    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
604    ///
605    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
606    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(
607    ///     vec![2, 3],
608    ///     vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0],
609    /// ).unwrap(), ctx.clone()).unwrap();
610    /// let y = x.reshape(&[6]).unwrap();
611    ///
612    /// assert_eq!(y.shape(), &[6]);
613    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
614    /// # Ok::<(), tenferro_ad::Error>(())
615    /// ```
616    ///
617    /// # Errors
618    ///
619    /// Returns [`tenferro_tensor::ValidationError::ShapeMismatch`] when the element count
620    /// changes, `InvalidArgument` when the target shape product overflows, or a
621    /// typed backend/runtime-state error.
622    pub fn reshape(&self, shape: &[usize]) -> Result<Self> {
623        let op = StdTensorOp::Reshape {
624            to_shape: DimExpr::from_concrete(shape),
625        };
626        // INVARIANT: a returned eager tensor cannot borrow `self`'s group, so
627        // the explicit duplicate precedes metadata-only view construction.
628        let base = self.to_tensor()?;
629        if let Ok(value) = TensorValue::from_tensor(base).reshape_view(shape) {
630            return Self::nary_value_op(&[self], op, value);
631        }
632        self.unary_op(op)
633    }
634
635    /// Slice with explicit start, limit, and stride per axis.
636    ///
637    /// # Examples
638    ///
639    /// ```
640    /// use tenferro_cpu::CpuBackend;
641    /// use tenferro_ad::{EagerRuntime, EagerTensor, SliceConfig, Tensor};
642    ///
643    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
644    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap(), ctx.clone()).unwrap();
645    /// let y = x
646    ///     .slice(SliceConfig {
647    ///         starts: vec![1],
648    ///         limits: vec![3],
649    ///         strides: vec![1],
650    ///     })
651    ///     .unwrap();
652    ///
653    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[2.0, 3.0]);
654    /// # Ok::<(), tenferro_ad::Error>(())
655    /// ```
656    ///
657    /// # Errors
658    ///
659    /// Returns [`tenferro_tensor::Error::Validation`] with
660    /// `AxisOutOfBounds`/`InvalidArgument` when starts, limits, or strides are
661    /// invalid, or a typed backend/runtime-state error while creating the view.
662    pub fn slice(&self, config: SliceConfig) -> Result<Self> {
663        // INVARIANT: the result must retain an independent owner while the
664        // input handle remains live; this is an explicit duplicate, not an
665        // implicit backend transfer.
666        let base = self.to_tensor()?;
667        let value = TensorValue::from_tensor(base)
668            .slice_view(&config)
669            .map_err(Error::TensorRuntime)?;
670        Self::nary_value_op(&[self], StdTensorOp::Slice(config), value)
671    }
672
673    /// Broadcast into a larger shape with explicit dimension placement.
674    ///
675    /// # Examples
676    ///
677    /// ```
678    /// use tenferro_cpu::CpuBackend;
679    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
680    ///
681    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
682    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx.clone()).unwrap();
683    /// let y = x.broadcast_in_dim(&[3, 2], &[0]).unwrap();
684    ///
685    /// assert_eq!(y.shape(), &[3, 2]);
686    /// # Ok::<(), tenferro_ad::Error>(())
687    /// ```
688    ///
689    /// # Errors
690    ///
691    /// Returns [`tenferro_tensor::Error::Validation`] with `AxisOutOfBounds`,
692    /// `DuplicateAxis`, or `ShapeMismatch` when `shape`/`dims` cannot broadcast
693    /// the input, or a typed backend/runtime-state error.
694    pub fn broadcast_in_dim(&self, shape: &[usize], dims: &[usize]) -> Result<Self> {
695        if let Some(error) = broadcast_in_dim_extent_error(self.shape(), shape, dims) {
696            return Err(broadcast_error("EagerTensor::broadcast_in_dim", error));
697        }
698        let op = StdTensorOp::BroadcastInDim {
699            shape: DimExpr::from_concrete(shape),
700            dims: dims.to_vec(),
701        };
702        // INVARIANT: output descriptors cannot borrow the input's move-only
703        // allocation group, so this explicit duplicate owns the view's root.
704        let base = self.to_tensor()?;
705        let value = TensorValue::from_tensor(base)
706            .broadcast_in_dim_view(shape, dims)
707            .map_err(Error::TensorRuntime)?;
708        Self::nary_value_op(&[self], op, value)
709    }
710
711    /// Convert the tensor to a different dtype using checked conversion.
712    ///
713    /// Use [`cast`](Self::cast) when a lossy dtype projection is intended.
714    ///
715    /// # Examples
716    ///
717    /// ```
718    /// use tenferro_cpu::CpuBackend;
719    /// use tenferro_ad::{DType, EagerRuntime, EagerTensor, Tensor};
720    ///
721    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
722    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, -2.0]).unwrap(), ctx.clone()).unwrap();
723    /// let y = x.convert(DType::C64).unwrap();
724    ///
725    /// assert_eq!(y.dtype(), DType::C64);
726    /// assert_eq!(y.shape(), &[2]);
727    /// # Ok::<(), tenferro_ad::Error>(())
728    /// ```
729    ///
730    /// # Errors
731    ///
732    /// Returns [`tenferro_tensor::Error::UnsupportedDTypeConversion`] when the
733    /// requested pair is outside tenferro's checked dtype-promotion lattice.
734    /// Use [`cast`](Self::cast) for explicit lossy projection; backend
735    /// execution can additionally return a typed runtime-state error.
736    pub fn convert(&self, to: DType) -> Result<Self> {
737        tenferro_tensor::validate::validate_convert_dtype("EagerTensor::convert", self.dtype(), to)
738            .map_err(Error::TensorRuntime)?;
739        self.cast(to)
740    }
741
742    /// Cast the tensor to a different dtype using explicit dtype projection.
743    ///
744    /// `cast` may truncate, narrow precision, project complex values to their
745    /// real component, or use boolean truthiness where the backend supports the
746    /// requested projection.
747    ///
748    /// # Examples
749    ///
750    /// ```
751    /// use tenferro_cpu::CpuBackend;
752    /// use tenferro_ad::{DType, EagerRuntime, EagerTensor, Tensor};
753    ///
754    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
755    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.2_f64, -2.8]).unwrap(), ctx.clone()).unwrap();
756    /// let y = x.cast(DType::I32).unwrap();
757    ///
758    /// assert_eq!(y.value().unwrap().as_slice::<i32>().unwrap(), &[1, -2]);
759    /// # Ok::<(), tenferro_ad::Error>(())
760    /// ```
761    /// # Errors
762    ///
763    /// Returns a typed [`tenferro_tensor::Error::Unsupported`] when the eager
764    /// backend cannot project the requested dtype, or a backend/runtime-state
765    /// error during execution.
766    pub fn cast(&self, to: DType) -> Result<Self> {
767        self.unary_op(StdTensorOp::Convert {
768            from: self.dtype(),
769            to,
770        })
771    }
772
773    /// Pad with zeros using StableHLO-style edge and interior padding.
774    ///
775    /// # Examples
776    ///
777    /// ```
778    /// use tenferro_cpu::CpuBackend;
779    /// use tenferro_ad::{EagerRuntime, EagerTensor, PadConfig, Tensor};
780    ///
781    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
782    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
783    /// let y = x
784    ///     .pad(PadConfig {
785    ///         edge_padding_low: vec![1],
786    ///         edge_padding_high: vec![1],
787    ///         interior_padding: vec![1],
788    ///     })
789    ///     .unwrap();
790    ///
791    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[0.0, 1.0, 0.0, 2.0, 0.0]);
792    /// # Ok::<(), tenferro_ad::Error>(())
793    /// ```
794    /// # Errors
795    ///
796    /// Returns [`tenferro_runtime::Error::TensorRuntime`] containing
797    /// [`tenferro_tensor::ValidationError::InvalidArgument`] when a
798    /// padding vector has a length different from the input rank, interior
799    /// padding is negative, or edge/interior padding produces a negative
800    /// dimension or checked output-size arithmetic overflows.
801    /// Backend execution and unavailable runtime state are propagated as their
802    /// typed [`tenferro_runtime::Error::TensorRuntime`] or
803    /// [`tenferro_runtime::Error::RuntimeState`] variants.
804    pub fn pad(&self, config: PadConfig) -> Result<Self> {
805        self.unary_op(StdTensorOp::Pad(config))
806    }
807
808    /// Reverse the order of elements along the requested axes.
809    ///
810    /// # Examples
811    ///
812    /// ```
813    /// use tenferro_cpu::CpuBackend;
814    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
815    ///
816    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
817    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap(), ctx.clone()).unwrap();
818    /// let y = x.reverse(&[0]).unwrap();
819    ///
820    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[4.0, 3.0, 2.0, 1.0]);
821    /// # Ok::<(), tenferro_ad::Error>(())
822    /// ```
823    /// # Errors
824    ///
825    /// Returns [`tenferro_tensor::Error::Validation`] with `AxisOutOfBounds` or
826    /// `DuplicateAxis` for an invalid axis list, or a typed backend/
827    /// runtime-state error during execution.
828    pub fn reverse(&self, axes: &[usize]) -> Result<Self> {
829        validate_eager_axes("EagerTensor::reverse", self.shape().len(), axes)?;
830        self.unary_op(StdTensorOp::Reverse {
831            axes: axes.to_vec(),
832        })
833    }
834
835    /// Gather slices from `self` using integer start indices.
836    ///
837    /// # Examples
838    ///
839    /// ```
840    /// use tenferro_cpu::CpuBackend;
841    /// use tenferro_ad::{EagerRuntime, EagerTensor, GatherConfig, Tensor};
842    ///
843    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
844    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(
845    ///     vec![5],
846    ///     vec![10.0_f64, 20.0, 30.0, 40.0, 50.0],
847    /// ).unwrap(), ctx.clone()).unwrap();
848    /// let indices = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![3], vec![4_i64, 1, 0]).unwrap(), ctx.clone()).unwrap();
849    /// let y = x
850    ///     .gather(
851    ///         &indices,
852    ///         GatherConfig {
853    ///             offset_dims: vec![],
854    ///             collapsed_slice_dims: vec![0],
855    ///             start_index_map: vec![0],
856    ///             index_vector_dim: 1,
857    ///             slice_sizes: vec![1],
858    ///         },
859    ///     )
860    ///     .unwrap();
861    ///
862    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[50.0, 20.0, 10.0]);
863    /// # Ok::<(), tenferro_ad::Error>(())
864    /// ```
865    /// # Errors
866    ///
867    /// Returns [`tenferro_tensor::Error::Validation`] when the gather
868    /// configuration has an invalid rank, axis, shape, or index dtype, or a
869    /// typed backend/runtime-state error.
870    pub fn gather(&self, indices: &Self, config: GatherConfig) -> Result<Self> {
871        self.binary_op(indices, StdTensorOp::Gather(config))
872    }
873
874    /// Scatter updates into `self` using StableHLO scatter semantics.
875    ///
876    /// # Examples
877    ///
878    /// ```
879    /// use tenferro_cpu::CpuBackend;
880    /// use tenferro_ad::{EagerRuntime, EagerTensor, ScatterConfig, Tensor};
881    ///
882    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
883    /// let operand = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![4], vec![0.0_f64, 0.0, 0.0, 0.0]).unwrap(), ctx.clone()).unwrap();
884    /// let indices = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2, 1], vec![1_i64, 3]).unwrap(), ctx.clone()).unwrap();
885    /// let updates = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![5.0_f64, 7.0]).unwrap(), ctx.clone()).unwrap();
886    /// let result = operand
887    ///     .scatter(
888    ///         &indices,
889    ///         &updates,
890    ///         ScatterConfig {
891    ///             update_window_dims: vec![],
892    ///             inserted_window_dims: vec![0],
893    ///             scatter_dims_to_operand_dims: vec![0],
894    ///             index_vector_dim: 1,
895    ///         },
896    ///     )
897    ///     .unwrap();
898    ///
899    /// assert_eq!(result.value().unwrap().as_slice::<f64>().unwrap(), &[0.0, 5.0, 0.0, 7.0]);
900    /// # Ok::<(), tenferro_ad::Error>(())
901    /// ```
902    /// # Errors
903    ///
904    /// Returns [`tenferro_tensor::Error::Validation`] when the scatter
905    /// configuration, index/update shapes, or index dtype is invalid, or a
906    /// typed backend/runtime-state error.
907    pub fn scatter(&self, indices: &Self, updates: &Self, config: ScatterConfig) -> Result<Self> {
908        self.ternary_op(indices, updates, StdTensorOp::Scatter(config))
909    }
910
911    /// Slice using runtime start indices.
912    ///
913    /// # Examples
914    ///
915    /// ```
916    /// use tenferro_cpu::CpuBackend;
917    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
918    ///
919    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
920    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![5], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0]).unwrap(), ctx.clone()).unwrap();
921    /// let starts = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2_i64]).unwrap(), ctx.clone()).unwrap();
922    /// let y = x.dynamic_slice(&starts, &[2]).unwrap();
923    ///
924    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[3.0, 4.0]);
925    /// # Ok::<(), tenferro_ad::Error>(())
926    /// ```
927    /// # Errors
928    ///
929    /// Returns [`tenferro_tensor::Error::Validation`] when `starts` has the
930    /// wrong dtype/shape or `sizes` exceeds the operand rank, including an
931    /// `AxisOutOfBounds` or `ShapeMismatch`, or a typed backend/runtime-state
932    /// error.
933    pub fn dynamic_slice(&self, starts: &Self, sizes: &[usize]) -> Result<Self> {
934        self.binary_op(
935            starts,
936            StdTensorOp::DynamicSlice {
937                slice_sizes: sizes.to_vec(),
938            },
939        )
940    }
941
942    /// Concatenate tensors along one axis.
943    ///
944    /// # Examples
945    ///
946    /// ```
947    /// use tenferro_cpu::CpuBackend;
948    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
949    ///
950    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
951    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
952    /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone()).unwrap();
953    /// let z = EagerTensor::concatenate(&[&x, &y], 0).unwrap();
954    ///
955    /// assert_eq!(z.value().unwrap().as_slice::<f64>().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
956    /// # Ok::<(), tenferro_ad::Error>(())
957    /// ```
958    /// # Errors
959    ///
960    /// Returns [`tenferro_tensor::ValidationError::InvalidArgument`] when `tensors` is
961    /// empty or `axis` is outside the rank, `ShapeMismatch`/`DTypeMismatch`
962    /// when inputs cannot be concatenated, or a typed backend/runtime-state
963    /// error.
964    pub fn concatenate(tensors: &[&Self], axis: usize) -> Result<Self> {
965        Self::nary_op(
966            tensors,
967            StdTensorOp::Concatenate {
968                axis,
969                input_count: tensors.len(),
970            },
971        )
972    }
973
974    /// Extract the diagonal along two axes.
975    ///
976    /// # Examples
977    ///
978    /// ```
979    /// use tenferro_cpu::CpuBackend;
980    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
981    ///
982    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
983    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(
984    ///     vec![3, 3],
985    ///     vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0],
986    /// ).unwrap(), ctx.clone()).unwrap();
987    /// let y = x.extract_diag(0, 1).unwrap();
988    ///
989    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[1.0, 5.0, 9.0]);
990    /// # Ok::<(), tenferro_ad::Error>(())
991    /// ```
992    /// # Errors
993    ///
994    /// Returns [`tenferro_tensor::Error::Validation`] with `RankMismatch`,
995    /// `AxisOutOfBounds`, or `DuplicateAxis` when the selected axes cannot form
996    /// a diagonal, or a typed backend/runtime-state error.
997    pub fn extract_diag(&self, axis_a: usize, axis_b: usize) -> Result<Self> {
998        self.unary_op(StdTensorOp::ExtractDiag { axis_a, axis_b })
999    }
1000
1001    /// Embed a vector or lower-rank tensor along a diagonal.
1002    ///
1003    /// # Examples
1004    ///
1005    /// ```
1006    /// use tenferro_cpu::CpuBackend;
1007    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1008    ///
1009    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1010    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx.clone()).unwrap();
1011    /// let y = x.embed_diag(0, 1).unwrap();
1012    ///
1013    /// assert_eq!(y.shape(), &[3, 3]);
1014    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0]);
1015    /// # Ok::<(), tenferro_ad::Error>(())
1016    /// ```
1017    /// # Errors
1018    ///
1019    /// Returns [`tenferro_tensor::Error::Validation`] with `RankMismatch`,
1020    /// `AxisOutOfBounds`, or `DuplicateAxis` when the diagonal axes are not
1021    /// valid for embedding, or a typed backend/runtime-state error.
1022    pub fn embed_diag(&self, axis_a: usize, axis_b: usize) -> Result<Self> {
1023        self.unary_op(StdTensorOp::EmbedDiag { axis_a, axis_b })
1024    }
1025
1026    /// Keep the lower triangle and zero the rest.
1027    ///
1028    /// # Examples
1029    ///
1030    /// ```
1031    /// use tenferro_cpu::CpuBackend;
1032    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1033    ///
1034    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1035    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap(), ctx.clone()).unwrap();
1036    /// let y = x.tril(0).unwrap();
1037    ///
1038    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[1.0, 2.0, 0.0, 4.0]);
1039    /// # Ok::<(), tenferro_ad::Error>(())
1040    /// ```
1041    /// # Errors
1042    ///
1043    /// Returns [`tenferro_tensor::ValidationError::RankMismatch`] when the operand is not
1044    /// a matrix, or a typed unsupported/backend/runtime-state error.
1045    pub fn tril(&self, k: i64) -> Result<Self> {
1046        self.unary_op(StdTensorOp::Tril { k })
1047    }
1048
1049    /// Keep the upper triangle and zero the rest.
1050    ///
1051    /// # Examples
1052    ///
1053    /// ```
1054    /// use tenferro_cpu::CpuBackend;
1055    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1056    ///
1057    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1058    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap(), ctx.clone()).unwrap();
1059    /// let y = x.triu(0).unwrap();
1060    ///
1061    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[1.0, 0.0, 3.0, 4.0]);
1062    /// # Ok::<(), tenferro_ad::Error>(())
1063    /// ```
1064    /// # Errors
1065    ///
1066    /// Returns [`tenferro_tensor::ValidationError::RankMismatch`] when the operand is not
1067    /// a matrix, or a typed unsupported/backend/runtime-state error.
1068    pub fn triu(&self, k: i64) -> Result<Self> {
1069        self.unary_op(StdTensorOp::Triu { k })
1070    }
1071
1072    /// Reduce product over the requested axes.
1073    ///
1074    /// # Examples
1075    ///
1076    /// ```
1077    /// use tenferro_cpu::CpuBackend;
1078    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1079    ///
1080    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1081    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap(), ctx.clone()).unwrap();
1082    /// let y = x.reduce_prod(None).unwrap();
1083    ///
1084    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[24.0]);
1085    /// # Ok::<(), tenferro_ad::Error>(())
1086    /// ```
1087    /// # Errors
1088    ///
1089    /// Returns [`tenferro_tensor::Error::Validation`] with `AxisOutOfBounds` or
1090    /// `DuplicateAxis` for an invalid reduction axis, or a typed
1091    /// unsupported/backend/runtime-state error for the selected dtype.
1092    pub fn reduce_prod(&self, axes: Option<&[usize]>) -> Result<Self> {
1093        let axes = axes.map_or_else(|| (0..self.shape().len()).collect(), <[usize]>::to_vec);
1094        validate_eager_axes("EagerTensor::reduce_prod", self.shape().len(), &axes)?;
1095        self.unary_op(StdTensorOp::ReduceProd { axes })
1096    }
1097
1098    /// Reduce maximum over the requested axes.
1099    ///
1100    /// # Examples
1101    ///
1102    /// ```
1103    /// use tenferro_cpu::CpuBackend;
1104    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1105    ///
1106    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1107    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap(), ctx.clone()).unwrap();
1108    /// let y = x.reduce_max(None).unwrap();
1109    ///
1110    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[4.0]);
1111    /// # Ok::<(), tenferro_ad::Error>(())
1112    /// ```
1113    /// # Errors
1114    ///
1115    /// Returns [`tenferro_tensor::Error::Validation`] with `AxisOutOfBounds` or
1116    /// `DuplicateAxis` for an invalid reduction axis, or a typed
1117    /// unsupported/backend/runtime-state error for the selected dtype.
1118    pub fn reduce_max(&self, axes: Option<&[usize]>) -> Result<Self> {
1119        let axes = axes.map_or_else(|| (0..self.shape().len()).collect(), <[usize]>::to_vec);
1120        validate_eager_axes("EagerTensor::reduce_max", self.shape().len(), &axes)?;
1121        self.unary_op(StdTensorOp::ReduceMax { axes })
1122    }
1123
1124    /// Reduce minimum over the requested axes.
1125    ///
1126    /// # Examples
1127    ///
1128    /// ```
1129    /// use tenferro_cpu::CpuBackend;
1130    /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1131    ///
1132    /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new())?;
1133    /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap(), ctx.clone()).unwrap();
1134    /// let y = x.reduce_min(None).unwrap();
1135    ///
1136    /// assert_eq!(y.value().unwrap().as_slice::<f64>().unwrap(), &[1.0]);
1137    /// # Ok::<(), tenferro_ad::Error>(())
1138    /// ```
1139    /// # Errors
1140    ///
1141    /// Returns [`tenferro_tensor::Error::Validation`] with `AxisOutOfBounds` or
1142    /// `DuplicateAxis` for an invalid reduction axis, or a typed
1143    /// unsupported/backend/runtime-state error for the selected dtype.
1144    pub fn reduce_min(&self, axes: Option<&[usize]>) -> Result<Self> {
1145        let axes = axes.map_or_else(|| (0..self.shape().len()).collect(), <[usize]>::to_vec);
1146        validate_eager_axes("EagerTensor::reduce_min", self.shape().len(), &axes)?;
1147        self.unary_op(StdTensorOp::ReduceMin { axes })
1148    }
1149
1150    pub(crate) fn unary_op(&self, op: StdTensorOp) -> Result<Self> {
1151        Self::nary_op(&[self], op)
1152    }
1153
1154    pub(crate) fn binary_op(&self, other: &Self, op: StdTensorOp) -> Result<Self> {
1155        Self::nary_op(&[self, other], op)
1156    }
1157
1158    pub(crate) fn ternary_op(&self, b: &Self, c: &Self, op: StdTensorOp) -> Result<Self> {
1159        Self::nary_op(&[self, b, c], op)
1160    }
1161
1162    pub(crate) fn nary_value_op(
1163        tensors: &[&Self],
1164        op: StdTensorOp,
1165        value: TensorValue,
1166    ) -> Result<Self> {
1167        let Some(first) = tensors.first() else {
1168            return Err(empty_nary_input_error(&op));
1169        };
1170
1171        let ctx = Arc::clone(&first.ctx);
1172        for tensor in tensors.iter().skip(1) {
1173            if !first.same_context(tensor) {
1174                return Err(Error::ContextMismatch {
1175                    lhs: first.ctx_id(),
1176                    rhs: tensor.ctx_id(),
1177                });
1178            }
1179        }
1180
1181        if !eager_grad_recording_enabled()
1182            || (!eager_capture_active() && !tensors.iter().any(|tensor| tensor.requires_grad))
1183        {
1184            return Self::new_untracked_value_result(ctx, value);
1185        }
1186
1187        let output_ref = &value;
1188        let mut recorded = record_eager_value_outputs(&op, &[output_ref], tensors)?;
1189        let trace = recorded.traces.pop().ok_or_else(|| {
1190            Error::Internal(format!("expected one eager trace for {:?}, got 0", op))
1191        })?;
1192        let semantic_trace = recorded.semantic_traces.pop().flatten();
1193
1194        Self::new_result_value(
1195            ctx,
1196            trace.key,
1197            value,
1198            trace.requires_grad,
1199            trace.trace,
1200            semantic_trace,
1201        )
1202    }
1203
1204    pub(crate) fn nary_op(tensors: &[&Self], op: StdTensorOp) -> Result<Self> {
1205        let total_started = eager_op_profile_start();
1206        let Some(first) = tensors.first() else {
1207            return Err(empty_nary_input_error(&op));
1208        };
1209        let expected = op.input_count();
1210        if tensors.len() != expected {
1211            return Err(wrong_nary_input_count_error(&op, expected, tensors.len()));
1212        }
1213
1214        let ctx = Arc::clone(&first.ctx);
1215        profile_eager_op_section("nary_op.context_check", || -> Result<()> {
1216            for tensor in tensors.iter().skip(1) {
1217                if !first.same_context(tensor) {
1218                    return Err(Error::ContextMismatch {
1219                        lhs: first.ctx_id(),
1220                        rhs: tensor.ctx_id(),
1221                    });
1222                }
1223            }
1224            Ok(())
1225        })?;
1226
1227        let any_requires_grad = profile_eager_op_section("nary_op.requires_grad_scan", || {
1228            eager_grad_recording_enabled()
1229                && (eager_capture_active() || tensors.iter().any(|tensor| tensor.requires_grad))
1230        });
1231        if !any_requires_grad {
1232            let input_reads = profile_eager_op_section("nary_op.collect_input_reads", || {
1233                tensors
1234                    .iter()
1235                    .map(|tensor| tensor.tensor_read())
1236                    .collect::<Vec<_>>()
1237            });
1238            let output = profile_eager_op_section("nary_op.exec_single_output_read", || {
1239                exec_single_output_read(&op, &input_reads, &ctx)
1240            })?;
1241            let result = profile_eager_op_section("nary_op.new_untracked_result", || {
1242                Self::new_untracked_result(ctx, output)
1243            });
1244            if let Some(total_started) = total_started {
1245                record_eager_op_profile("nary_op.total", total_started.elapsed());
1246                maybe_print_eager_op_profile();
1247            }
1248            return result;
1249        }
1250
1251        let input_arcs = profile_eager_op_section("nary_op.materialize_inputs", || {
1252            tensors
1253                .iter()
1254                .map(|tensor| tensor.to_tensor().map(Arc::new))
1255                .collect::<Result<Vec<_>>>()
1256        })?;
1257        let inputs: Vec<&Tensor> = profile_eager_op_section("nary_op.collect_inputs", || {
1258            input_arcs.iter().map(|tensor| tensor.as_ref()).collect()
1259        });
1260        let output = profile_eager_op_section("nary_op.exec_single_output", || {
1261            exec_single_output(&op, &inputs, &ctx)
1262        })?;
1263
1264        let outputs = vec![&output];
1265        let mut recorded = profile_eager_op_section("nary_op.record_outputs", || {
1266            record_eager_outputs(&op, &outputs, tensors)
1267        })?;
1268        let trace = recorded.traces.pop().ok_or_else(|| {
1269            Error::Internal(format!("expected one eager trace for {:?}, got 0", op))
1270        })?;
1271        let semantic_trace = recorded.semantic_traces.pop().flatten();
1272
1273        let result = profile_eager_op_section("nary_op.new_tracked_result", || {
1274            Self::new_result_with_semantic_trace(
1275                ctx,
1276                trace.key,
1277                output,
1278                trace.requires_grad,
1279                trace.trace,
1280                semantic_trace,
1281            )
1282        });
1283        if let Some(total_started) = total_started {
1284            record_eager_op_profile("nary_op.total", total_started.elapsed());
1285            maybe_print_eager_op_profile();
1286        }
1287        result
1288    }
1289}
1290
1291fn validate_eager_axes(op: &'static str, rank: usize, axes: &[usize]) -> Result<()> {
1292    tenferro_tensor::validate::validate_unique_axes(op, "axis", rank, axes)
1293        .map_err(Error::TensorRuntime)
1294}
1295
1296fn validate_eager_dot_general_config(
1297    _op: &'static str,
1298    config: &DotGeneralConfig,
1299    lhs_rank: usize,
1300    rhs_rank: usize,
1301) -> Result<()> {
1302    config
1303        .validate_dims_with_ranks(lhs_rank, rhs_rank)
1304        .map_err(Error::TensorRuntime)
1305}
1306
1307fn empty_nary_input_error(op: &StdTensorOp) -> Error {
1308    Error::TensorRuntime(tenferro_tensor::Error::invalid_argument(
1309        eager_validation_op_name(op),
1310        "inputs",
1311        "operation requires at least one input tensor",
1312    ))
1313}
1314
1315fn wrong_nary_input_count_error(op: &StdTensorOp, expected: usize, actual: usize) -> Error {
1316    Error::TensorRuntime(tenferro_tensor::Error::invalid_argument(
1317        eager_validation_op_name(op),
1318        "inputs",
1319        format!("operation expects {expected} inputs, got {actual}"),
1320    ))
1321}
1322
1323fn eager_validation_op_name(op: &StdTensorOp) -> &'static str {
1324    match op {
1325        StdTensorOp::Concatenate { .. } => "concatenate",
1326        _ => "eager_nary_op",
1327    }
1328}
1329
1330fn finite_real_factor(value: f64) -> Result<f64> {
1331    if value.is_finite() {
1332        Ok(value)
1333    } else {
1334        Err(Error::TensorRuntime(
1335            tenferro_tensor::Error::invalid_argument(
1336                "scale_real",
1337                "factor",
1338                format!("real scalar must be finite, got {value}"),
1339            ),
1340        ))
1341    }
1342}
1343
1344fn round_real_to_i64(value: f64) -> Result<i64> {
1345    let rounded = finite_real_factor(value)?.round();
1346    if rounded < i64::MIN as f64 || rounded >= -(i64::MIN as f64) {
1347        return Err(Error::TensorRuntime(
1348            tenferro_tensor::Error::invalid_argument(
1349                "scale_real",
1350                "factor",
1351                format!("rounded real scalar {rounded} is out of i64 range"),
1352            ),
1353        ));
1354    }
1355    Ok(rounded as i64)
1356}
1357
1358fn round_real_to_i32(value: f64) -> Result<i32> {
1359    let rounded = round_real_to_i64(value)?;
1360    i32::try_from(rounded).map_err(|_| {
1361        Error::TensorRuntime(tenferro_tensor::Error::invalid_argument(
1362            "scale_real",
1363            "factor",
1364            format!("rounded real scalar {rounded} is out of i32 range"),
1365        ))
1366    })
1367}
1368
1369fn bool_from_real(value: f64) -> Result<bool> {
1370    Ok(finite_real_factor(value)? != 0.0)
1371}
tenferro_ad/eager_ops.rs

tenferro_ad/
eager_ops.rs
