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