tenferro_ad/eager/composite.rs
1//! Eager composite operations (activations, normalizations, softmax,
2//! `take_along_axis`) on the borrowed [`EagerSession`].
3//!
4//! The formulation and edge-case policy are shared with the traced and
5//! concrete-session surfaces through `tenferro_runtime::composite`; AD follows
6//! from the recorded primitives.
7
8use tenferro_runtime::composite::{
9 self, scalar_tensor, zero_pad_config, CompositeBinary, CompositeOps, CompositeReduce,
10 CompositeUnary,
11};
12use tenferro_tensor::{ActivationOp, CompareDir, DType, GatherConfig};
13
14use super::{EagerSession, EagerTensor};
15use crate::{Error, Result};
16
17/// [`CompositeOps`] over a borrowed eager session.
18struct EagerComposite<'s, 'a> {
19 session: &'s mut EagerSession<'a>,
20}
21
22impl CompositeOps for EagerComposite<'_, '_> {
23 type Value = EagerTensor;
24 type Error = Error;
25
26 fn dtype(&self, value: &EagerTensor) -> DType {
27 value.dtype()
28 }
29
30 fn shape(&self, value: &EagerTensor) -> Result<Vec<usize>> {
31 Ok(value.shape().to_vec())
32 }
33
34 fn scalar(&mut self, dtype: DType, value: f64) -> Result<EagerTensor> {
35 self.session
36 .constant_from_host(scalar_tensor(dtype, value)?)
37 }
38
39 fn unary(&mut self, op: CompositeUnary, value: &EagerTensor) -> Result<EagerTensor> {
40 let session = &mut *self.session;
41 match op {
42 CompositeUnary::Neg => session.neg(value),
43 CompositeUnary::Exp => session.exp(value),
44 CompositeUnary::Log => session.log(value),
45 CompositeUnary::Log1p => session.log1p(value),
46 CompositeUnary::Tanh => session.tanh(value),
47 CompositeUnary::Erf => session.erf(value),
48 CompositeUnary::Rsqrt => session.rsqrt(value),
49 }
50 }
51
52 fn binary(
53 &mut self,
54 op: CompositeBinary,
55 lhs: &EagerTensor,
56 rhs: &EagerTensor,
57 ) -> Result<EagerTensor> {
58 let session = &mut *self.session;
59 match op {
60 CompositeBinary::Add => session.add(lhs, rhs),
61 CompositeBinary::Sub => session.sub(lhs, rhs),
62 CompositeBinary::Mul => session.mul(lhs, rhs),
63 CompositeBinary::Div => session.div(lhs, rhs),
64 CompositeBinary::Maximum => session.maximum(lhs, rhs),
65 }
66 }
67
68 fn compare(
69 &mut self,
70 lhs: &EagerTensor,
71 rhs: &EagerTensor,
72 dir: CompareDir,
73 ) -> Result<EagerTensor> {
74 self.session.compare(lhs, rhs, dir)
75 }
76
77 fn select(
78 &mut self,
79 condition: &EagerTensor,
80 on_true: &EagerTensor,
81 on_false: &EagerTensor,
82 ) -> Result<EagerTensor> {
83 self.session.where_select(condition, on_true, on_false)
84 }
85
86 fn reduce(
87 &mut self,
88 op: CompositeReduce,
89 value: &EagerTensor,
90 axes: &[usize],
91 ) -> Result<EagerTensor> {
92 match op {
93 CompositeReduce::Sum => self.session.reduce_sum(value, Some(axes)),
94 CompositeReduce::Max => self.session.reduce_max(value, Some(axes)),
95 CompositeReduce::SumSquares => self.session.reduce_sum_squares(value, Some(axes)),
96 }
97 }
98
99 fn broadcast_in_dim(
100 &mut self,
101 value: &EagerTensor,
102 shape: &[usize],
103 dims: &[usize],
104 ) -> Result<EagerTensor> {
105 self.session.broadcast_in_dim(value, shape, dims)
106 }
107
108 fn reshape(&mut self, value: &EagerTensor, shape: &[usize]) -> Result<EagerTensor> {
109 self.session.reshape(value, shape.to_vec())
110 }
111
112 fn concatenate(&mut self, values: &[&EagerTensor], axis: usize) -> Result<EagerTensor> {
113 self.session.concatenate(values, axis)
114 }
115
116 fn pad(&mut self, value: &EagerTensor, low: &[usize], high: &[usize]) -> Result<EagerTensor> {
117 self.session.pad(value, zero_pad_config(low, high))
118 }
119
120 fn gather(
121 &mut self,
122 operand: &EagerTensor,
123 indices: &EagerTensor,
124 config: GatherConfig,
125 ) -> Result<EagerTensor> {
126 self.session.gather(operand, indices, config)
127 }
128}
129
130impl<'a> EagerSession<'a> {
131 fn composite(&mut self) -> EagerComposite<'_, 'a> {
132 EagerComposite { session: self }
133 }
134
135 /// Run a fused activation through the backend when no AD record is needed.
136 ///
137 /// Returns `Ok(None)` when the operands are tracked, semantic capture is
138 /// active, or the backend declines the fused form; the caller then records
139 /// the shared composite formulation instead.
140 fn fused_activation(
141 &mut self,
142 op: ActivationOp,
143 input: &EagerTensor,
144 ) -> Result<Option<EagerTensor>> {
145 if !crate::eager_ops::untracked_fast_path_allowed(&[input]) {
146 return Ok(None);
147 }
148 let read = input.tensor_read();
149 let Some(output) = self.engine.fused_activation_read(op, read)? else {
150 return Ok(None);
151 };
152 Ok(Some(EagerTensor::new_untracked_result(
153 std::sync::Arc::clone(&input.ctx),
154 output,
155 )?))
156 }
157
158 /// Logistic sigmoid `1 / (1 + exp(-x))`, overflow-free.
159 ///
160 /// Evaluated as `1 / (1 + e)` for `x > 0` and `e / (1 + e)` otherwise, with
161 /// `e = exp(-|x|)`; the derivative is finite everywhere (`sigmoid'(0) = 1/4`).
162 /// Real `F32`/`F64` only.
163 ///
164 /// # Examples
165 /// ```rust
166 /// use tenferro_ad::{EagerRuntime, Tensor};
167 /// let ctx = EagerRuntime::new()?;
168 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
169 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![3], vec![-700.0_f64, 0.0, 1000.0])?)?;
170 /// s.sigmoid(&x)
171 /// })?;
172 /// let y = y.value()?;
173 /// let y = y.as_slice::<f64>()?;
174 /// assert_eq!(y[1], 0.5);
175 /// assert!(y[0] > 0.0 && y[0] < 1e-300);
176 /// assert_eq!(y[2], 1.0);
177 /// # Ok::<(), tenferro_ad::Error>(())
178 /// ```
179 ///
180 /// # Errors
181 ///
182 /// Returns a typed `UnsupportedDType` error for complex, integer, or `Bool`
183 /// input, [`Error::ContextMismatch`] for a tensor from another runtime, or a
184 /// backend error.
185 pub fn sigmoid(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
186 if let Some(out) = self.fused_activation(ActivationOp::Sigmoid, input)? {
187 return Ok(out);
188 }
189 composite::sigmoid(&mut self.composite(), input)
190 }
191
192 /// SiLU (swish) `x * sigmoid(x)`.
193 ///
194 /// Real `F32`/`F64` only.
195 ///
196 /// # Examples
197 /// ```rust
198 /// use tenferro_ad::{EagerRuntime, Tensor};
199 /// let ctx = EagerRuntime::new()?;
200 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
201 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![3], vec![-1.0_f64, 0.0, 1.0])?)?;
202 /// s.silu(&x)
203 /// })?;
204 /// let y = y.value()?;
205 /// let y = y.as_slice::<f64>()?;
206 /// assert_eq!(y[1], 0.0);
207 /// assert!((y[2] - 1.0 / (1.0 + (-1.0_f64).exp())).abs() < 1e-15);
208 /// # Ok::<(), tenferro_ad::Error>(())
209 /// ```
210 ///
211 /// # Errors
212 ///
213 /// Returns a typed `UnsupportedDType` error for complex, integer, or `Bool`
214 /// input, [`Error::ContextMismatch`] for a tensor from another runtime, or a
215 /// backend error.
216 pub fn silu(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
217 if let Some(out) = self.fused_activation(ActivationOp::Silu, input)? {
218 return Ok(out);
219 }
220 composite::silu(&mut self.composite(), input)
221 }
222
223 /// Softplus `log(1 + exp(x))` in the stable form `max(x, 0) + log1p(exp(-|x|))`.
224 ///
225 /// Never overflows; `softplus'(0) = 1/2` and `softplus''(0) = 1/4`. Real `F32`/`F64` only.
226 ///
227 /// # Examples
228 /// ```rust
229 /// use tenferro_ad::{EagerRuntime, Tensor};
230 /// let ctx = EagerRuntime::new()?;
231 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
232 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![3], vec![-1000.0_f64, 0.0, 1000.0])?)?;
233 /// s.softplus(&x)
234 /// })?;
235 /// let y = y.value()?;
236 /// let y = y.as_slice::<f64>()?;
237 /// assert_eq!(y[0], 0.0);
238 /// assert!((y[1] - 2.0_f64.ln()).abs() < 1e-15);
239 /// assert_eq!(y[2], 1000.0);
240 /// # Ok::<(), tenferro_ad::Error>(())
241 /// ```
242 ///
243 /// # Errors
244 ///
245 /// Returns a typed `UnsupportedDType` error for complex, integer, or `Bool`
246 /// input, [`Error::ContextMismatch`] for a tensor from another runtime, or a
247 /// backend error.
248 pub fn softplus(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
249 if let Some(out) = self.fused_activation(ActivationOp::Softplus, input)? {
250 return Ok(out);
251 }
252 composite::softplus(&mut self.composite(), input)
253 }
254
255 /// Exact GELU `x/2 * (1 + erf(x / sqrt(2)))` (PyTorch `approximate="none"`).
256 ///
257 /// Real `F32`/`F64` only.
258 ///
259 /// # Examples
260 /// ```rust
261 /// use tenferro_ad::{EagerRuntime, Tensor};
262 /// let ctx = EagerRuntime::new()?;
263 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
264 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![3], vec![-1.0_f64, 0.0, 1.0])?)?;
265 /// s.gelu(&x)
266 /// })?;
267 /// let y = y.value()?;
268 /// let y = y.as_slice::<f64>()?;
269 /// assert_eq!(y[1], 0.0);
270 /// assert!((y[2] - 0.841_344_746_068_542_9).abs() < 1e-15);
271 /// # Ok::<(), tenferro_ad::Error>(())
272 /// ```
273 ///
274 /// # Errors
275 ///
276 /// Returns a typed `UnsupportedDType` error for complex, integer, or `Bool`
277 /// input, [`Error::ContextMismatch`] for a tensor from another runtime, or a
278 /// backend error.
279 pub fn gelu(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
280 if let Some(out) = self.fused_activation(ActivationOp::Gelu, input)? {
281 return Ok(out);
282 }
283 composite::gelu(&mut self.composite(), input)
284 }
285
286 /// GELU tanh approximation (PyTorch `approximate="tanh"`).
287 ///
288 /// `x/2 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 x^3)))`; real `F32`/`F64` only.
289 ///
290 /// # Examples
291 /// ```rust
292 /// use tenferro_ad::{EagerRuntime, Tensor};
293 /// let ctx = EagerRuntime::new()?;
294 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
295 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![3], vec![-1.0_f64, 0.0, 1.0])?)?;
296 /// s.gelu_tanh(&x)
297 /// })?;
298 /// let y = y.value()?;
299 /// let y = y.as_slice::<f64>()?;
300 /// assert_eq!(y[1], 0.0);
301 /// assert!((y[2] - 0.841_191_990_608_276_8).abs() < 1e-12);
302 /// # Ok::<(), tenferro_ad::Error>(())
303 /// ```
304 ///
305 /// # Errors
306 ///
307 /// Returns a typed `UnsupportedDType` error for complex, integer, or `Bool`
308 /// input, [`Error::ContextMismatch`] for a tensor from another runtime, or a
309 /// backend error.
310 pub fn gelu_tanh(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
311 if let Some(out) = self.fused_activation(ActivationOp::GeluTanh, input)? {
312 return Ok(out);
313 }
314 composite::gelu_tanh(&mut self.composite(), input)
315 }
316
317 /// Arithmetic mean over `axes` (`None` reduces every axis).
318 ///
319 /// Float and complex dtypes. The sum is divided by the element count; a mean
320 /// over zero elements is `NaN`, and `Some(&[])` is the identity.
321 ///
322 /// # Examples
323 /// ```rust
324 /// use tenferro_ad::{EagerRuntime, Tensor};
325 /// let ctx = EagerRuntime::new()?;
326 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
327 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
328 /// s.reduce_mean(&x, Some(&[1]))
329 /// })?;
330 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[2.0, 3.0]);
331 /// # Ok::<(), tenferro_ad::Error>(())
332 /// ```
333 ///
334 /// # Errors
335 ///
336 /// Returns a typed `UnsupportedDType` error for integer or `Bool` input, an
337 /// `AxisOutOfBounds` / `DuplicateAxis` validation error for invalid axes,
338 /// [`Error::ContextMismatch`] for a tensor from another runtime, or a backend error.
339 pub fn reduce_mean(
340 &mut self,
341 input: &EagerTensor,
342 axes: Option<&[usize]>,
343 ) -> Result<EagerTensor> {
344 composite::reduce_mean(&mut self.composite(), input, axes)
345 }
346
347 /// Max-subtracted softmax along `axis`.
348 ///
349 /// A slice that is entirely `-inf` returns zeros with a finite gradient instead
350 /// of `NaN`; a participating `NaN` or `+inf` makes its slice `NaN`; a
351 /// zero-length `axis` returns an empty result. Real `F32`/`F64` only.
352 ///
353 /// # Examples
354 /// ```rust
355 /// use tenferro_ad::{EagerRuntime, Tensor};
356 /// let ctx = EagerRuntime::new()?;
357 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
358 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![2], vec![0.0_f64, f64::NEG_INFINITY])?)?;
359 /// s.softmax(&x, 0)
360 /// })?;
361 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[1.0, 0.0]);
362 /// # Ok::<(), tenferro_ad::Error>(())
363 /// ```
364 ///
365 /// # Errors
366 ///
367 /// Returns a typed `UnsupportedDType` error for non-real input, an
368 /// `AxisOutOfBounds` validation error for an invalid axis,
369 /// [`Error::ContextMismatch`] for a tensor from another runtime, or a backend error.
370 pub fn softmax(&mut self, input: &EagerTensor, axis: usize) -> Result<EagerTensor> {
371 if crate::eager_ops::untracked_fast_path_allowed(&[input])
372 && let Some(output) = self.engine.softmax_read(input.tensor_read(), axis)?
373 {
374 return EagerTensor::new_untracked_result(std::sync::Arc::clone(&input.ctx), output);
375 }
376 composite::softmax(&mut self.composite(), input, axis)
377 }
378
379 /// Max-subtracted log-softmax along `axis`.
380 ///
381 /// A slice that is entirely `-inf` returns `-inf` with a finite gradient
382 /// instead of `NaN`. Real `F32`/`F64` only.
383 ///
384 /// # Examples
385 /// ```rust
386 /// use tenferro_ad::{EagerRuntime, Tensor};
387 /// let ctx = EagerRuntime::new()?;
388 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
389 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 1.0])?)?;
390 /// s.log_softmax(&x, 0)
391 /// })?;
392 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[-std::f64::consts::LN_2; 2]);
393 /// # Ok::<(), tenferro_ad::Error>(())
394 /// ```
395 ///
396 /// # Errors
397 ///
398 /// Returns a typed `UnsupportedDType` error for non-real input, an
399 /// `AxisOutOfBounds` validation error for an invalid axis,
400 /// [`Error::ContextMismatch`] for a tensor from another runtime, or a backend error.
401 pub fn log_softmax(&mut self, input: &EagerTensor, axis: usize) -> Result<EagerTensor> {
402 composite::log_softmax(&mut self.composite(), input, axis)
403 }
404
405 /// Softmax along `axis` over the entries where the `Bool` `mask` is true.
406 ///
407 /// `mask` broadcasts to the input shape. Masked-out entries get probability `0`
408 /// and a zero gradient whatever their value; a slice with no unmasked entry
409 /// returns zeros with a zero gradient.
410 ///
411 /// # Examples
412 /// ```rust
413 /// use tenferro_ad::{EagerRuntime, Tensor};
414 /// let ctx = EagerRuntime::new()?;
415 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
416 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 1.0, f64::NAN])?)?;
417 /// let mask = s.constant_from_host(Tensor::from_vec_col_major(vec![3], vec![true, true, false])?)?;
418 /// s.masked_softmax(&x, &mask, 0)
419 /// })?;
420 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.5, 0.5, 0.0]);
421 /// # Ok::<(), tenferro_ad::Error>(())
422 /// ```
423 ///
424 /// # Errors
425 ///
426 /// Returns a typed `UnsupportedDType` error for non-real input, a
427 /// `DTypeMismatch` validation error for a non-`Bool` mask, `ShapeMismatch` for
428 /// a mask that does not broadcast to the input, `AxisOutOfBounds` for an invalid
429 /// axis, [`Error::ContextMismatch`] for a tensor from another runtime, or a
430 /// backend error.
431 pub fn masked_softmax(
432 &mut self,
433 input: &EagerTensor,
434 mask: &EagerTensor,
435 axis: usize,
436 ) -> Result<EagerTensor> {
437 composite::masked_softmax(&mut self.composite(), input, mask, axis)
438 }
439
440 /// Log-softmax along `axis` over the entries where the `Bool` `mask` is true.
441 ///
442 /// Masked-out entries are `-inf` with a zero gradient; a slice with no unmasked
443 /// entry is all `-inf` with a zero gradient.
444 ///
445 /// # Examples
446 /// ```rust
447 /// use tenferro_ad::{EagerRuntime, Tensor};
448 /// let ctx = EagerRuntime::new()?;
449 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
450 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0])?)?;
451 /// let mask = s.constant_from_host(Tensor::from_vec_col_major(vec![2], vec![true, false])?)?;
452 /// s.masked_log_softmax(&x, &mask, 0)
453 /// })?;
454 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0, f64::NEG_INFINITY]);
455 /// # Ok::<(), tenferro_ad::Error>(())
456 /// ```
457 ///
458 /// # Errors
459 ///
460 /// Returns a typed `UnsupportedDType` error for non-real input, a
461 /// `DTypeMismatch` validation error for a non-`Bool` mask, `ShapeMismatch` for
462 /// a mask that does not broadcast to the input, `AxisOutOfBounds` for an invalid
463 /// axis, [`Error::ContextMismatch`] for a tensor from another runtime, or a
464 /// backend error.
465 pub fn masked_log_softmax(
466 &mut self,
467 input: &EagerTensor,
468 mask: &EagerTensor,
469 axis: usize,
470 ) -> Result<EagerTensor> {
471 composite::masked_log_softmax(&mut self.composite(), input, mask, axis)
472 }
473
474 /// Layer normalization along `axis` with optional affine `weight` / `bias`.
475 ///
476 /// `(x - mean) / sqrt(var + eps) * weight + bias` with the biased variance of the
477 /// centered values; `weight` and `bias` are rank-1 of length `shape[axis]`. A
478 /// zero-variance slice normalizes to `0` (then `bias`) with a finite gradient
479 /// when `eps > 0`. Real `F32`/`F64` only.
480 ///
481 /// # Examples
482 /// ```rust
483 /// use tenferro_ad::{EagerRuntime, Tensor};
484 /// let ctx = EagerRuntime::new()?;
485 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
486 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 3.0])?)?;
487 /// s.layer_norm(&x, 0, None, None, 0.0)
488 /// })?;
489 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[-1.0, 1.0]);
490 /// # Ok::<(), tenferro_ad::Error>(())
491 /// ```
492 ///
493 /// # Errors
494 ///
495 /// Returns a typed `UnsupportedDType` error for non-real input, an
496 /// `AxisOutOfBounds` validation error for an invalid axis, `InvalidArgument` for
497 /// a negative or non-finite `eps`, `DTypeMismatch` / `ShapeMismatch` for a weight
498 /// or bias that is not a same-dtype vector of the axis length,
499 /// [`Error::ContextMismatch`] for a tensor from another runtime, or a backend error.
500 pub fn layer_norm(
501 &mut self,
502 input: &EagerTensor,
503 axis: usize,
504 weight: Option<&EagerTensor>,
505 bias: Option<&EagerTensor>,
506 eps: f64,
507 ) -> Result<EagerTensor> {
508 composite::layer_norm(&mut self.composite(), input, axis, weight, bias, eps)
509 }
510
511 /// RMS normalization along `axis` with optional affine `weight` / `bias`.
512 ///
513 /// `x / sqrt(mean(x^2) + eps) * weight + bias`; `weight` and `bias` are rank-1 of
514 /// length `shape[axis]`. An all-zero slice normalizes to `0` (then `bias`) with
515 /// a finite gradient when `eps > 0`. Real `F32`/`F64` only.
516 ///
517 /// # Examples
518 /// ```rust
519 /// use tenferro_ad::{EagerRuntime, Tensor};
520 /// let ctx = EagerRuntime::new()?;
521 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
522 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![2], vec![0.0_f64, 0.0])?)?;
523 /// s.rms_norm(&x, 0, None, None, 1e-6)
524 /// })?;
525 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0, 0.0]);
526 /// # Ok::<(), tenferro_ad::Error>(())
527 /// ```
528 ///
529 /// # Errors
530 ///
531 /// Returns a typed `UnsupportedDType` error for non-real input, an
532 /// `AxisOutOfBounds` validation error for an invalid axis, `InvalidArgument` for
533 /// a negative or non-finite `eps`, `DTypeMismatch` / `ShapeMismatch` for a weight
534 /// or bias that is not a same-dtype vector of the axis length,
535 /// [`Error::ContextMismatch`] for a tensor from another runtime, or a backend error.
536 pub fn rms_norm(
537 &mut self,
538 input: &EagerTensor,
539 axis: usize,
540 weight: Option<&EagerTensor>,
541 bias: Option<&EagerTensor>,
542 eps: f64,
543 ) -> Result<EagerTensor> {
544 composite::rms_norm(&mut self.composite(), input, axis, weight, bias, eps)
545 }
546
547 /// NumPy-style `take_along_axis` over `gather`.
548 ///
549 /// `out[.., i, ..] = input[.., indices[.., i, ..], ..]` along `axis`. `indices`
550 /// (I32/I64) has the input's rank; every other dimension is either the input's
551 /// extent (batch-varying indices) or `1` (the whole extent is taken). Indices
552 /// must be in bounds. The gradient flows to `input` only.
553 ///
554 /// # Examples
555 /// ```rust
556 /// use tenferro_ad::{EagerRuntime, Tensor};
557 /// let ctx = EagerRuntime::new()?;
558 /// // Per-batch row gather: out[i, j, b] = x[idx[i, b], j, b].
559 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
560 /// let x = s.constant_from_host(Tensor::from_vec_col_major(vec![2, 2, 2], (0..8).map(f64::from).collect::<Vec<_>>())?)?;
561 /// let idx = s.constant_from_host(Tensor::from_vec_col_major(vec![2, 1, 2], vec![1_i64, 0, 0, 0])?)?;
562 /// s.take_along_axis(&x, &idx, 0)
563 /// })?;
564 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[1.0, 0.0, 3.0, 2.0, 4.0, 4.0, 6.0, 6.0]);
565 /// # Ok::<(), tenferro_ad::Error>(())
566 /// ```
567 ///
568 /// # Errors
569 ///
570 /// Returns a `RankMismatch` / `ShapeMismatch` validation error for incompatible
571 /// index shapes, `AxisOutOfBounds` for an invalid axis, `InvalidArgument` when
572 /// taking from a zero-length axis, a typed `UnsupportedDType` error for a
573 /// non-integer index dtype, [`Error::ContextMismatch`] for a tensor from another
574 /// runtime, or a backend error.
575 pub fn take_along_axis(
576 &mut self,
577 input: &EagerTensor,
578 indices: &EagerTensor,
579 axis: usize,
580 ) -> Result<EagerTensor> {
581 composite::take_along_axis(&mut self.composite(), input, indices, axis)
582 }
583}