tenferro_ad/eager_ops.rs
1use std::sync::Arc;
2
3use computegraph::GraphOperation;
4use tenferro_ops::broadcast::{
5 broadcast_input_plan, broadcast_shape, broadcast_shapes, BroadcastError,
6};
7use tenferro_ops::dim_expr::DimExpr;
8use tenferro_ops::std_tensor_op::StdTensorOp;
9use tenferro_tensor::{
10 DType, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig, Tensor,
11 TensorValue,
12};
13
14use crate::eager::{
15 eager_grad_recording_enabled, exec_single_output, exec_single_output_read,
16 maybe_print_eager_op_profile, profile_eager_op_section, record_eager_op_profile,
17 record_eager_outputs, EagerTensor,
18};
19use crate::eager_exec::exec_dot_general_with_conj_on_tensor_reads;
20use crate::error::{Error, Result};
21use crate::metadata::push_metadata_scope;
22
23pub(crate) fn broadcast_binary(
24 op: &'static str,
25 lhs: &EagerTensor,
26 rhs: &EagerTensor,
27) -> Result<(EagerTensor, EagerTensor)> {
28 ensure_same_context(lhs, rhs)?;
29 let shape =
30 broadcast_shape(lhs.shape(), rhs.shape()).map_err(|err| broadcast_error(op, err))?;
31 Ok((
32 broadcast_to(op, lhs, &shape)?,
33 broadcast_to(op, rhs, &shape)?,
34 ))
35}
36
37pub(crate) fn broadcast_ternary(
38 op: &'static str,
39 first: &EagerTensor,
40 second: &EagerTensor,
41 third: &EagerTensor,
42) -> Result<(EagerTensor, EagerTensor, EagerTensor)> {
43 ensure_same_context(first, second)?;
44 ensure_same_context(first, third)?;
45 let shape = broadcast_shapes([first.shape(), second.shape(), third.shape()])
46 .map_err(|err| broadcast_error(op, err))?;
47 Ok((
48 broadcast_to(op, first, &shape)?,
49 broadcast_to(op, second, &shape)?,
50 broadcast_to(op, third, &shape)?,
51 ))
52}
53
54fn broadcast_to(
55 op: &'static str,
56 input: &EagerTensor,
57 target_shape: &[usize],
58) -> Result<EagerTensor> {
59 let input_shape = input.shape();
60 if input_shape == target_shape {
61 return Ok(input.clone());
62 }
63
64 let plan =
65 broadcast_input_plan(input_shape, target_shape).map_err(|err| broadcast_error(op, err))?;
66 let source = if plan.source_shape == input_shape {
67 input.clone()
68 } else {
69 input.reshape(&plan.source_shape)?
70 };
71 source.broadcast_in_dim(target_shape, &plan.dims)
72}
73
74fn broadcast_error(op: &'static str, err: BroadcastError) -> Error {
75 match err {
76 BroadcastError::IncompatibleBinary { lhs, rhs } => {
77 tenferro_tensor::Error::ShapeMismatch { op, lhs, rhs }.into()
78 }
79 BroadcastError::IncompatibleInput { input, output }
80 | BroadcastError::RankTooLarge { input, output } => tenferro_tensor::Error::InvalidConfig {
81 op,
82 message: format!("cannot broadcast shape {input:?} to {output:?}"),
83 }
84 .into(),
85 }
86}
87
88fn ensure_same_context(lhs: &EagerTensor, rhs: &EagerTensor) -> Result<()> {
89 if !lhs.same_context(rhs) {
90 return Err(Error::ContextMismatch {
91 lhs: lhs.ctx_id(),
92 rhs: rhs.ctx_id(),
93 });
94 }
95 Ok(())
96}
97
98impl std::ops::Add for &EagerTensor {
99 type Output = Result<EagerTensor>;
100
101 fn add(self, rhs: &EagerTensor) -> Result<EagerTensor> {
102 EagerTensor::add(self, rhs)
103 }
104}
105
106impl std::ops::Sub for &EagerTensor {
107 type Output = Result<EagerTensor>;
108
109 fn sub(self, rhs: &EagerTensor) -> Result<EagerTensor> {
110 EagerTensor::sub(self, rhs)
111 }
112}
113
114impl std::ops::Mul for &EagerTensor {
115 type Output = Result<EagerTensor>;
116
117 fn mul(self, rhs: &EagerTensor) -> Result<EagerTensor> {
118 EagerTensor::mul(self, rhs)
119 }
120}
121
122impl std::ops::Div for &EagerTensor {
123 type Output = Result<EagerTensor>;
124
125 fn div(self, rhs: &EagerTensor) -> Result<EagerTensor> {
126 EagerTensor::div(self, rhs)
127 }
128}
129
130impl std::ops::Rem for &EagerTensor {
131 type Output = Result<EagerTensor>;
132
133 fn rem(self, rhs: &EagerTensor) -> Result<EagerTensor> {
134 EagerTensor::rem(self, rhs)
135 }
136}
137
138impl std::ops::Neg for &EagerTensor {
139 type Output = Result<EagerTensor>;
140
141 fn neg(self) -> Result<EagerTensor> {
142 EagerTensor::neg(self)
143 }
144}
145
146impl EagerTensor {
147 /// Elementwise addition.
148 ///
149 /// # Examples
150 ///
151 /// ```
152 /// use tenferro_cpu::CpuBackend;
153 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
154 ///
155 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
156 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
157 /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone()).unwrap();
158 /// let z = x.add(&y).unwrap();
159 ///
160 /// assert_eq!(z.materialized().unwrap().as_slice::<f64>().unwrap(), &[4.0, 6.0]);
161 /// ```
162 pub fn add(&self, other: &Self) -> Result<Self> {
163 let (lhs, rhs) = broadcast_binary("add", self, other)?;
164 lhs.binary_op(&rhs, StdTensorOp::Add)
165 }
166
167 /// Elementwise subtraction.
168 pub fn sub(&self, other: &Self) -> Result<Self> {
169 let (lhs, rhs) = broadcast_binary("sub", self, other)?;
170 lhs.binary_op(&rhs, StdTensorOp::Sub)
171 }
172
173 /// Elementwise multiplication.
174 ///
175 /// # Examples
176 ///
177 /// ```
178 /// use tenferro_cpu::CpuBackend;
179 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
180 ///
181 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
182 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
183 /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone()).unwrap();
184 /// let z = x.mul(&y).unwrap();
185 ///
186 /// assert_eq!(z.materialized().unwrap().as_slice::<f64>().unwrap(), &[3.0, 8.0]);
187 /// ```
188 pub fn mul(&self, other: &Self) -> Result<Self> {
189 let (lhs, rhs) = broadcast_binary("mul", self, other)?;
190 lhs.binary_op(&rhs, StdTensorOp::Mul)
191 }
192
193 /// Negate the tensor.
194 ///
195 /// # Examples
196 ///
197 /// ```
198 /// use tenferro_cpu::CpuBackend;
199 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
200 ///
201 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
202 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, -2.0]).unwrap(), ctx.clone()).unwrap();
203 /// let y = x.neg().unwrap();
204 ///
205 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[-1.0, 2.0]);
206 /// ```
207 pub fn neg(&self) -> Result<Self> {
208 self.unary_op(StdTensorOp::Neg)
209 }
210
211 /// Elementwise exponential.
212 ///
213 /// # Examples
214 ///
215 /// ```
216 /// use tenferro_cpu::CpuBackend;
217 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
218 ///
219 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
220 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![0.0_f64]).unwrap(), ctx.clone()).unwrap();
221 /// let y = x.exp().unwrap();
222 ///
223 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[1.0]);
224 /// ```
225 pub fn exp(&self) -> Result<Self> {
226 self.unary_op(StdTensorOp::Exp)
227 }
228
229 /// Reduce sum over the requested axes.
230 ///
231 /// # Examples
232 ///
233 /// ```
234 /// use tenferro_cpu::CpuBackend;
235 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
236 ///
237 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
238 /// 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();
239 /// let y = x.reduce_sum(&[0, 1]).unwrap();
240 ///
241 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[10.0]);
242 /// ```
243 pub fn reduce_sum(&self, axes: &[usize]) -> Result<Self> {
244 validate_eager_axes("EagerTensor::reduce_sum", self.shape().len(), axes)?;
245 self.unary_op(StdTensorOp::ReduceSum {
246 axes: axes.to_vec(),
247 })
248 }
249
250 /// Execute a dot-general contraction eagerly.
251 ///
252 /// # Examples
253 ///
254 /// ```
255 /// use tenferro_cpu::CpuBackend;
256 /// use tenferro_ad::{DotGeneralConfig, EagerRuntime, EagerTensor, Tensor};
257 ///
258 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
259 /// 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();
260 /// 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();
261 /// let c = a.dot_general(&b, DotGeneralConfig {
262 /// lhs_contracting_dims: vec![1],
263 /// rhs_contracting_dims: vec![0],
264 /// lhs_batch_dims: vec![],
265 /// rhs_batch_dims: vec![],
266 /// }).unwrap();
267 ///
268 /// assert_eq!(c.shape(), &[2, 2]);
269 /// ```
270 pub fn dot_general(&self, other: &Self, config: DotGeneralConfig) -> Result<Self> {
271 validate_eager_dot_general_config(
272 "EagerTensor::dot_general",
273 &config,
274 self.shape().len(),
275 other.shape().len(),
276 )?;
277 self.binary_op(other, StdTensorOp::DotGeneral { config })
278 }
279
280 /// Execute a dot-general contraction, optionally conjugating either operand.
281 ///
282 /// Untracked tensors route the conjugation flags directly to the backend so
283 /// the conjugated operand does not need to be materialized. Tracked tensors
284 /// fall back to explicit `Conj` plus `DotGeneral` so reverse-mode AD keeps
285 /// the same graph semantics as the standard eager ops.
286 pub fn dot_general_with_conj(
287 &self,
288 other: &Self,
289 config: &DotGeneralConfig,
290 lhs_conj: bool,
291 rhs_conj: bool,
292 ) -> Result<Self> {
293 if !self.same_context(other) {
294 return Err(Error::ContextMismatch {
295 lhs: self.ctx_id(),
296 rhs: other.ctx_id(),
297 });
298 }
299 validate_eager_dot_general_config(
300 "EagerTensor::dot_general_with_conj",
301 config,
302 self.shape().len(),
303 other.shape().len(),
304 )?;
305
306 if !self.requires_grad && !other.requires_grad {
307 let ctx = Arc::clone(&self.ctx);
308 let output = ctx.with_backend_mut(|backend| {
309 exec_dot_general_with_conj_on_tensor_reads(
310 self.tensor_read(),
311 other.tensor_read(),
312 config,
313 lhs_conj,
314 rhs_conj,
315 backend,
316 )
317 })??;
318 return Self::new_untracked_result(ctx, output);
319 }
320
321 match (lhs_conj, rhs_conj) {
322 (false, false) => self.dot_general(other, config.clone()),
323 (true, false) => self.conj()?.dot_general(other, config.clone()),
324 (false, true) => {
325 let rhs = other.conj()?;
326 self.dot_general(&rhs, config.clone())
327 }
328 (true, true) => {
329 let lhs = self.conj()?;
330 let rhs = other.conj()?;
331 lhs.dot_general(&rhs, config.clone())
332 }
333 }
334 }
335
336 /// Matrix multiplication for rank-2 tensors.
337 ///
338 /// This is a convenience wrapper over [`Self::dot_general`] that
339 /// contracts the left matrix's column axis with the right matrix's row
340 /// axis.
341 ///
342 /// # Examples
343 ///
344 /// ```
345 /// use tenferro_cpu::CpuBackend;
346 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
347 ///
348 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
349 /// let a = EagerTensor::from_tensor_in(
350 /// Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap(),
351 /// ctx.clone(),
352 /// ).unwrap();
353 /// let b = EagerTensor::from_tensor_in(
354 /// Tensor::from_vec_col_major(vec![2, 1], vec![5.0_f64, 6.0]).unwrap(),
355 /// ctx,
356 /// ).unwrap();
357 /// let c = a.matmul(&b).unwrap();
358 ///
359 /// assert_eq!(c.shape(), &[2, 1]);
360 /// assert_eq!(c.materialized().unwrap().as_slice::<f64>().unwrap(), &[23.0, 34.0]);
361 /// ```
362 pub fn matmul(&self, other: &Self) -> Result<Self> {
363 let lhs_shape = self.shape();
364 let rhs_shape = other.shape();
365 if lhs_shape.len() != 2 {
366 return Err(tenferro_tensor::Error::RankMismatch {
367 op: "matmul",
368 expected: 2,
369 actual: lhs_shape.len(),
370 }
371 .into());
372 }
373 if rhs_shape.len() != 2 {
374 return Err(tenferro_tensor::Error::RankMismatch {
375 op: "matmul",
376 expected: 2,
377 actual: rhs_shape.len(),
378 }
379 .into());
380 }
381 if lhs_shape[1] != rhs_shape[0] {
382 return Err(tenferro_tensor::Error::ShapeMismatch {
383 op: "matmul",
384 lhs: lhs_shape.to_vec(),
385 rhs: rhs_shape.to_vec(),
386 }
387 .into());
388 }
389 self.dot_general(
390 other,
391 DotGeneralConfig {
392 lhs_contracting_dims: vec![1],
393 rhs_contracting_dims: vec![0],
394 lhs_batch_dims: vec![],
395 rhs_batch_dims: vec![],
396 },
397 )
398 }
399
400 /// Permute tensor axes.
401 ///
402 /// # Examples
403 ///
404 /// ```
405 /// use tenferro_cpu::CpuBackend;
406 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
407 ///
408 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
409 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(
410 /// vec![2, 3],
411 /// vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0],
412 /// ).unwrap(), ctx.clone()).unwrap();
413 /// let y = x.transpose(&[1, 0]).unwrap();
414 ///
415 /// assert_eq!(y.shape(), &[3, 2]);
416 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[1.0, 3.0, 5.0, 2.0, 4.0, 6.0]);
417 /// ```
418 pub fn transpose(&self, perm: &[usize]) -> Result<Self> {
419 let op = StdTensorOp::Transpose {
420 perm: perm.to_vec(),
421 };
422 let value = self
423 .value
424 .transpose_view(perm)
425 .map_err(Error::TensorRuntime)?;
426 Self::nary_value_op(&[self], op, value)
427 }
428
429 /// Reshape without changing element order.
430 ///
431 /// # Examples
432 ///
433 /// ```
434 /// use tenferro_cpu::CpuBackend;
435 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
436 ///
437 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
438 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(
439 /// vec![2, 3],
440 /// vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0],
441 /// ).unwrap(), ctx.clone()).unwrap();
442 /// let y = x.reshape(&[6]).unwrap();
443 ///
444 /// assert_eq!(y.shape(), &[6]);
445 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
446 /// ```
447 pub fn reshape(&self, shape: &[usize]) -> Result<Self> {
448 let op = StdTensorOp::Reshape {
449 to_shape: DimExpr::from_concrete(shape),
450 };
451 if let Ok(value) = self.value.reshape_view(shape) {
452 return Self::nary_value_op(&[self], op, value);
453 }
454 self.unary_op(op)
455 }
456
457 /// Slice with explicit start, limit, and stride per axis.
458 ///
459 /// # Examples
460 ///
461 /// ```
462 /// use tenferro_cpu::CpuBackend;
463 /// use tenferro_ad::{EagerRuntime, EagerTensor, SliceConfig, Tensor};
464 ///
465 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
466 /// 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();
467 /// let y = x
468 /// .slice(SliceConfig {
469 /// starts: vec![1],
470 /// limits: vec![3],
471 /// strides: vec![1],
472 /// })
473 /// .unwrap();
474 ///
475 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[2.0, 3.0]);
476 /// ```
477 pub fn slice(&self, config: SliceConfig) -> Result<Self> {
478 let value = self
479 .value
480 .slice_view(&config)
481 .map_err(Error::TensorRuntime)?;
482 Self::nary_value_op(&[self], StdTensorOp::Slice(config), value)
483 }
484
485 /// Broadcast into a larger shape with explicit dimension placement.
486 ///
487 /// # Examples
488 ///
489 /// ```
490 /// use tenferro_cpu::CpuBackend;
491 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
492 ///
493 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
494 /// 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();
495 /// let y = x.broadcast_in_dim(&[3, 2], &[0]).unwrap();
496 ///
497 /// assert_eq!(y.shape(), &[3, 2]);
498 /// ```
499 pub fn broadcast_in_dim(&self, shape: &[usize], dims: &[usize]) -> Result<Self> {
500 let op = StdTensorOp::BroadcastInDim {
501 shape: DimExpr::from_concrete(shape),
502 dims: dims.to_vec(),
503 };
504 let value = self
505 .value
506 .broadcast_in_dim_view(shape, dims)
507 .map_err(Error::TensorRuntime)?;
508 Self::nary_value_op(&[self], op, value)
509 }
510
511 /// Convert the tensor to a different dtype using checked conversion.
512 ///
513 /// Use [`cast`](Self::cast) when a lossy dtype projection is intended.
514 ///
515 /// # Examples
516 ///
517 /// ```
518 /// use tenferro_cpu::CpuBackend;
519 /// use tenferro_ad::{DType, EagerRuntime, EagerTensor, Tensor};
520 ///
521 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
522 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, -2.0]).unwrap(), ctx.clone()).unwrap();
523 /// let y = x.convert(DType::C64).unwrap();
524 ///
525 /// assert_eq!(y.dtype(), DType::C64);
526 /// assert_eq!(y.shape(), &[2]);
527 /// ```
528 ///
529 /// # Errors
530 ///
531 /// Returns an error when the requested conversion is outside tenferro's
532 /// checked dtype-promotion lattice. Use [`cast`](Self::cast) for explicit
533 /// lossy dtype projection.
534 pub fn convert(&self, to: DType) -> Result<Self> {
535 tenferro_tensor::validate::validate_convert_dtype("EagerTensor::convert", self.dtype(), to)
536 .map_err(Error::TensorRuntime)?;
537 self.cast(to)
538 }
539
540 /// Cast the tensor to a different dtype using explicit dtype projection.
541 ///
542 /// `cast` may truncate, narrow precision, project complex values to their
543 /// real component, or use boolean truthiness where the backend supports the
544 /// requested projection.
545 ///
546 /// # Examples
547 ///
548 /// ```
549 /// use tenferro_cpu::CpuBackend;
550 /// use tenferro_ad::{DType, EagerRuntime, EagerTensor, Tensor};
551 ///
552 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
553 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.2_f64, -2.8]).unwrap(), ctx.clone()).unwrap();
554 /// let y = x.cast(DType::I32).unwrap();
555 ///
556 /// assert_eq!(y.materialized().unwrap().as_slice::<i32>().unwrap(), &[1, -2]);
557 /// ```
558 pub fn cast(&self, to: DType) -> Result<Self> {
559 self.unary_op(StdTensorOp::Convert {
560 from: self.dtype(),
561 to,
562 })
563 }
564
565 /// Pad with zeros using StableHLO-style edge and interior padding.
566 ///
567 /// # Examples
568 ///
569 /// ```
570 /// use tenferro_cpu::CpuBackend;
571 /// use tenferro_ad::{EagerRuntime, EagerTensor, PadConfig, Tensor};
572 ///
573 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
574 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
575 /// let y = x
576 /// .pad(PadConfig {
577 /// edge_padding_low: vec![1],
578 /// edge_padding_high: vec![1],
579 /// interior_padding: vec![1],
580 /// })
581 /// .unwrap();
582 ///
583 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[0.0, 1.0, 0.0, 2.0, 0.0]);
584 /// ```
585 pub fn pad(&self, config: PadConfig) -> Result<Self> {
586 self.unary_op(StdTensorOp::Pad(config))
587 }
588
589 /// Reverse the order of elements along the requested axes.
590 ///
591 /// # Examples
592 ///
593 /// ```
594 /// use tenferro_cpu::CpuBackend;
595 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
596 ///
597 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
598 /// 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();
599 /// let y = x.reverse(&[0]).unwrap();
600 ///
601 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[4.0, 3.0, 2.0, 1.0]);
602 /// ```
603 pub fn reverse(&self, axes: &[usize]) -> Result<Self> {
604 validate_eager_axes("EagerTensor::reverse", self.shape().len(), axes)?;
605 self.unary_op(StdTensorOp::Reverse {
606 axes: axes.to_vec(),
607 })
608 }
609
610 /// Gather slices from `self` using integer start indices.
611 ///
612 /// # Examples
613 ///
614 /// ```
615 /// use tenferro_cpu::CpuBackend;
616 /// use tenferro_ad::{EagerRuntime, EagerTensor, GatherConfig, Tensor};
617 ///
618 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
619 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(
620 /// vec![5],
621 /// vec![10.0_f64, 20.0, 30.0, 40.0, 50.0],
622 /// ).unwrap(), ctx.clone()).unwrap();
623 /// let indices = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![3], vec![4_i64, 1, 0]).unwrap(), ctx.clone()).unwrap();
624 /// let y = x
625 /// .gather(
626 /// &indices,
627 /// GatherConfig {
628 /// offset_dims: vec![],
629 /// collapsed_slice_dims: vec![0],
630 /// start_index_map: vec![0],
631 /// index_vector_dim: 1,
632 /// slice_sizes: vec![1],
633 /// },
634 /// )
635 /// .unwrap();
636 ///
637 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[50.0, 20.0, 10.0]);
638 /// ```
639 pub fn gather(&self, indices: &Self, config: GatherConfig) -> Result<Self> {
640 self.binary_op(indices, StdTensorOp::Gather(config))
641 }
642
643 /// Scatter updates into `self` using StableHLO scatter semantics.
644 ///
645 /// # Examples
646 ///
647 /// ```
648 /// use tenferro_cpu::CpuBackend;
649 /// use tenferro_ad::{EagerRuntime, EagerTensor, ScatterConfig, Tensor};
650 ///
651 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
652 /// 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();
653 /// let indices = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2, 1], vec![1_i64, 3]).unwrap(), ctx.clone()).unwrap();
654 /// let updates = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![5.0_f64, 7.0]).unwrap(), ctx.clone()).unwrap();
655 /// let result = operand
656 /// .scatter(
657 /// &indices,
658 /// &updates,
659 /// ScatterConfig {
660 /// update_window_dims: vec![],
661 /// inserted_window_dims: vec![0],
662 /// scatter_dims_to_operand_dims: vec![0],
663 /// index_vector_dim: 1,
664 /// },
665 /// )
666 /// .unwrap();
667 ///
668 /// assert_eq!(result.materialized().unwrap().as_slice::<f64>().unwrap(), &[0.0, 5.0, 0.0, 7.0]);
669 /// ```
670 pub fn scatter(&self, indices: &Self, updates: &Self, config: ScatterConfig) -> Result<Self> {
671 self.ternary_op(indices, updates, StdTensorOp::Scatter(config))
672 }
673
674 /// Slice using runtime start indices.
675 ///
676 /// # Examples
677 ///
678 /// ```
679 /// use tenferro_cpu::CpuBackend;
680 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
681 ///
682 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
683 /// 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();
684 /// let starts = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2_i64]).unwrap(), ctx.clone()).unwrap();
685 /// let y = x.dynamic_slice(&starts, &[2]).unwrap();
686 ///
687 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[3.0, 4.0]);
688 /// ```
689 pub fn dynamic_slice(&self, starts: &Self, sizes: &[usize]) -> Result<Self> {
690 self.binary_op(
691 starts,
692 StdTensorOp::DynamicSlice {
693 slice_sizes: sizes.to_vec(),
694 },
695 )
696 }
697
698 /// Concatenate tensors along one axis.
699 ///
700 /// # Examples
701 ///
702 /// ```
703 /// use tenferro_cpu::CpuBackend;
704 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
705 ///
706 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
707 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
708 /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone()).unwrap();
709 /// let z = EagerTensor::concatenate(&[&x, &y], 0).unwrap();
710 ///
711 /// assert_eq!(z.materialized().unwrap().as_slice::<f64>().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
712 /// ```
713 pub fn concatenate(tensors: &[&Self], axis: usize) -> Result<Self> {
714 Self::nary_op(
715 tensors,
716 StdTensorOp::Concatenate {
717 axis,
718 input_count: tensors.len(),
719 },
720 )
721 }
722
723 /// Extract the diagonal along two axes.
724 ///
725 /// # Examples
726 ///
727 /// ```
728 /// use tenferro_cpu::CpuBackend;
729 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
730 ///
731 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
732 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(
733 /// vec![3, 3],
734 /// vec![1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0],
735 /// ).unwrap(), ctx.clone()).unwrap();
736 /// let y = x.extract_diag(0, 1).unwrap();
737 ///
738 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[1.0, 5.0, 9.0]);
739 /// ```
740 pub fn extract_diag(&self, axis_a: usize, axis_b: usize) -> Result<Self> {
741 self.unary_op(StdTensorOp::ExtractDiag { axis_a, axis_b })
742 }
743
744 /// Embed a vector or lower-rank tensor along a diagonal.
745 ///
746 /// # Examples
747 ///
748 /// ```
749 /// use tenferro_cpu::CpuBackend;
750 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
751 ///
752 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
753 /// 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();
754 /// let y = x.embed_diag(0, 1).unwrap();
755 ///
756 /// assert_eq!(y.shape(), &[3, 3]);
757 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0]);
758 /// ```
759 pub fn embed_diag(&self, axis_a: usize, axis_b: usize) -> Result<Self> {
760 self.unary_op(StdTensorOp::EmbedDiag { axis_a, axis_b })
761 }
762
763 /// Keep the lower triangle and zero the rest.
764 ///
765 /// # Examples
766 ///
767 /// ```
768 /// use tenferro_cpu::CpuBackend;
769 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
770 ///
771 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
772 /// 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();
773 /// let y = x.tril(0).unwrap();
774 ///
775 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[1.0, 2.0, 0.0, 4.0]);
776 /// ```
777 pub fn tril(&self, k: i64) -> Result<Self> {
778 self.unary_op(StdTensorOp::Tril { k })
779 }
780
781 /// Keep the upper triangle and zero the rest.
782 ///
783 /// # Examples
784 ///
785 /// ```
786 /// use tenferro_cpu::CpuBackend;
787 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
788 ///
789 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
790 /// 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();
791 /// let y = x.triu(0).unwrap();
792 ///
793 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[1.0, 0.0, 3.0, 4.0]);
794 /// ```
795 pub fn triu(&self, k: i64) -> Result<Self> {
796 self.unary_op(StdTensorOp::Triu { k })
797 }
798
799 /// Reduce product over the requested axes.
800 ///
801 /// # Examples
802 ///
803 /// ```
804 /// use tenferro_cpu::CpuBackend;
805 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
806 ///
807 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
808 /// 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();
809 /// let y = x.reduce_prod(&[0, 1]).unwrap();
810 ///
811 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[24.0]);
812 /// ```
813 pub fn reduce_prod(&self, axes: &[usize]) -> Result<Self> {
814 validate_eager_axes("EagerTensor::reduce_prod", self.shape().len(), axes)?;
815 self.unary_op(StdTensorOp::ReduceProd {
816 axes: axes.to_vec(),
817 })
818 }
819
820 /// Reduce maximum over the requested axes.
821 ///
822 /// # Examples
823 ///
824 /// ```
825 /// use tenferro_cpu::CpuBackend;
826 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
827 ///
828 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
829 /// 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();
830 /// let y = x.reduce_max(&[0, 1]).unwrap();
831 ///
832 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[4.0]);
833 /// ```
834 pub fn reduce_max(&self, axes: &[usize]) -> Result<Self> {
835 validate_eager_axes("EagerTensor::reduce_max", self.shape().len(), axes)?;
836 self.unary_op(StdTensorOp::ReduceMax {
837 axes: axes.to_vec(),
838 })
839 }
840
841 /// Reduce minimum over the requested axes.
842 ///
843 /// # Examples
844 ///
845 /// ```
846 /// use tenferro_cpu::CpuBackend;
847 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
848 ///
849 /// let ctx = EagerRuntime::with_cpu_backend(CpuBackend::new());
850 /// 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();
851 /// let y = x.reduce_min(&[0, 1]).unwrap();
852 ///
853 /// assert_eq!(y.materialized().unwrap().as_slice::<f64>().unwrap(), &[1.0]);
854 /// ```
855 pub fn reduce_min(&self, axes: &[usize]) -> Result<Self> {
856 validate_eager_axes("EagerTensor::reduce_min", self.shape().len(), axes)?;
857 self.unary_op(StdTensorOp::ReduceMin {
858 axes: axes.to_vec(),
859 })
860 }
861
862 pub(crate) fn unary_op(&self, op: StdTensorOp) -> Result<Self> {
863 Self::nary_op(&[self], op)
864 }
865
866 pub(crate) fn binary_op(&self, other: &Self, op: StdTensorOp) -> Result<Self> {
867 Self::nary_op(&[self, other], op)
868 }
869
870 pub(crate) fn ternary_op(&self, b: &Self, c: &Self, op: StdTensorOp) -> Result<Self> {
871 Self::nary_op(&[self, b, c], op)
872 }
873
874 pub(crate) fn nary_value_op(
875 tensors: &[&Self],
876 op: StdTensorOp,
877 value: TensorValue,
878 ) -> Result<Self> {
879 let Some(first) = tensors.first() else {
880 return Err(empty_nary_input_error(&op));
881 };
882
883 let ctx = Arc::clone(&first.ctx);
884 for tensor in tensors.iter().skip(1) {
885 if !first.same_context(tensor) {
886 return Err(Error::ContextMismatch {
887 lhs: first.ctx_id(),
888 rhs: tensor.ctx_id(),
889 });
890 }
891 }
892
893 if !eager_grad_recording_enabled() || !tensors.iter().any(|tensor| tensor.requires_grad) {
894 return Ok(Self::new_untracked_value_result(ctx, value));
895 }
896
897 let output = Arc::new(value.to_tensor().map_err(Error::from)?);
898 let outputs = vec![Arc::clone(&output)];
899 let mut recorded = record_eager_outputs(&op, &outputs, tensors)?;
900 let trace = recorded.traces.pop().ok_or_else(|| {
901 Error::Internal(format!("expected one eager trace for {:?}, got 0", op))
902 })?;
903 let mut metadata_scopes = vec![Arc::clone(&recorded.metadata_scope)];
904 for tensor in tensors {
905 for scope in &tensor.metadata_scopes {
906 push_metadata_scope(&mut metadata_scopes, Arc::clone(scope));
907 }
908 }
909
910 Self::new_result_value(
911 ctx,
912 trace.key,
913 value,
914 trace.requires_grad,
915 trace.trace,
916 metadata_scopes,
917 )
918 }
919
920 pub(crate) fn nary_op(tensors: &[&Self], op: StdTensorOp) -> Result<Self> {
921 let total_started = std::time::Instant::now();
922 let Some(first) = tensors.first() else {
923 return Err(empty_nary_input_error(&op));
924 };
925 let expected = op.input_count();
926 if tensors.len() != expected {
927 return Err(wrong_nary_input_count_error(&op, expected, tensors.len()));
928 }
929
930 let ctx = Arc::clone(&first.ctx);
931 profile_eager_op_section("nary_op.context_check", || -> Result<()> {
932 for tensor in tensors.iter().skip(1) {
933 if !first.same_context(tensor) {
934 return Err(Error::ContextMismatch {
935 lhs: first.ctx_id(),
936 rhs: tensor.ctx_id(),
937 });
938 }
939 }
940 Ok(())
941 })?;
942
943 let any_requires_grad = profile_eager_op_section("nary_op.requires_grad_scan", || {
944 eager_grad_recording_enabled() && tensors.iter().any(|tensor| tensor.requires_grad)
945 });
946 if !any_requires_grad {
947 let input_reads = profile_eager_op_section("nary_op.collect_input_reads", || {
948 tensors
949 .iter()
950 .map(|tensor| tensor.tensor_read())
951 .collect::<Vec<_>>()
952 });
953 let output = profile_eager_op_section("nary_op.exec_single_output_read", || {
954 exec_single_output_read(&op, &input_reads, &ctx)
955 })?;
956 let result = profile_eager_op_section("nary_op.new_untracked_result", || {
957 Self::new_untracked_result(ctx, output)
958 });
959 record_eager_op_profile("nary_op.total", total_started.elapsed());
960 maybe_print_eager_op_profile();
961 return result;
962 }
963
964 let input_arcs = profile_eager_op_section("nary_op.materialize_inputs", || {
965 tensors
966 .iter()
967 .map(|tensor| tensor.materialized_arc())
968 .collect::<Result<Vec<_>>>()
969 })?;
970 let inputs: Vec<&Tensor> = profile_eager_op_section("nary_op.collect_inputs", || {
971 input_arcs.iter().map(|tensor| tensor.as_ref()).collect()
972 });
973 let output = profile_eager_op_section("nary_op.exec_single_output", || {
974 exec_single_output(&op, &inputs, &ctx)
975 })?;
976
977 let output = Arc::new(output);
978 let outputs = vec![Arc::clone(&output)];
979 let mut recorded = profile_eager_op_section("nary_op.record_outputs", || {
980 record_eager_outputs(&op, &outputs, tensors)
981 })?;
982 let trace = recorded.traces.pop().ok_or_else(|| {
983 Error::Internal(format!("expected one eager trace for {:?}, got 0", op))
984 })?;
985 let mut metadata_scopes = vec![Arc::clone(&recorded.metadata_scope)];
986 for tensor in tensors {
987 for scope in &tensor.metadata_scopes {
988 push_metadata_scope(&mut metadata_scopes, Arc::clone(scope));
989 }
990 }
991
992 let result = profile_eager_op_section("nary_op.new_tracked_result", || {
993 Self::new_result_arc(
994 ctx,
995 trace.key,
996 output,
997 trace.requires_grad,
998 trace.trace,
999 metadata_scopes,
1000 )
1001 });
1002 record_eager_op_profile("nary_op.total", total_started.elapsed());
1003 maybe_print_eager_op_profile();
1004 result
1005 }
1006}
1007
1008fn validate_eager_axes(op: &'static str, rank: usize, axes: &[usize]) -> Result<()> {
1009 tenferro_tensor::validate::validate_unique_axes(op, "axis", rank, axes)
1010 .map_err(Error::TensorRuntime)
1011}
1012
1013fn validate_eager_dot_general_config(
1014 op: &'static str,
1015 config: &DotGeneralConfig,
1016 lhs_rank: usize,
1017 rhs_rank: usize,
1018) -> Result<()> {
1019 config
1020 .validate_dims_with_ranks(lhs_rank, rhs_rank)
1021 .map_err(|err| {
1022 Error::TensorRuntime(tenferro_tensor::Error::InvalidConfig {
1023 op,
1024 message: err.to_string(),
1025 })
1026 })
1027}
1028
1029fn empty_nary_input_error(op: &StdTensorOp) -> Error {
1030 Error::TensorRuntime(tenferro_tensor::Error::InvalidConfig {
1031 op: eager_validation_op_name(op),
1032 message: "operation requires at least one input tensor".to_string(),
1033 })
1034}
1035
1036fn wrong_nary_input_count_error(op: &StdTensorOp, expected: usize, actual: usize) -> Error {
1037 Error::TensorRuntime(tenferro_tensor::Error::InvalidConfig {
1038 op: eager_validation_op_name(op),
1039 message: format!("operation expects {expected} inputs, got {actual}"),
1040 })
1041}
1042
1043fn eager_validation_op_name(op: &StdTensorOp) -> &'static str {
1044 match op {
1045 StdTensorOp::Concatenate { .. } => "concatenate",
1046 _ => "eager_nary_op",
1047 }
1048}