665 lines
20 KiB
Rust
665 lines
20 KiB
Rust
//! Automatic differentiation engine for RustyTorch++
|
|
//!
|
|
//! This crate provides the `Autodiff<B>` decorator pattern for automatic
|
|
//! differentiation, inspired by Burn's architecture.
|
|
//!
|
|
//! ## Key Features
|
|
//!
|
|
//! - **Zero inference overhead**: Use raw backend without `Autodiff` wrapper
|
|
//! - **No global state**: Per-tensor gradient nodes with `Arc` reference counting
|
|
//! - **Composable**: `Autodiff<Quantized<CudaBackend>>` is possible
|
|
//! - **Type-safe**: Training vs inference is a type-level distinction
|
|
//!
|
|
//! ## Usage
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_backend_cuda::CudaBackend;
|
|
//! use rtx_autograd::{Autodiff, AutodiffBackend, no_grad};
|
|
//!
|
|
//! // Training with gradients
|
|
//! type TrainingBackend = Autodiff<CudaBackend>;
|
|
//! let x = TrainingBackend::from_data(&[1.0, 2.0], [2], &device);
|
|
//! let x = TrainingBackend::require_grad(x);
|
|
//! let y = TrainingBackend::mul(&x, &x);
|
|
//! let grads = TrainingBackend::backward(&y);
|
|
//!
|
|
//! // Inference without gradients (zero overhead!)
|
|
//! type InferenceBackend = CudaBackend;
|
|
//! let z = InferenceBackend::mul(&a, &b); // No gradient tracking!
|
|
//!
|
|
//! // Disable gradients temporarily
|
|
//! let result = no_grad(|| {
|
|
//! TrainingBackend::mul(&x, &y) // No gradient node created
|
|
//! });
|
|
//! ```
|
|
//!
|
|
//! ## Architecture
|
|
//!
|
|
//! ```text
|
|
//! Autodiff<B: Backend>
|
|
//! |
|
|
//! +-- AutodiffTensor<B, D> (wrapper around B::TensorPrimitive<D>)
|
|
//! | |
|
|
//! | +-- inner: B::TensorPrimitive<D>
|
|
//! | +-- node: Option<Arc<AutodiffNode<B>>>
|
|
//! | +-- id: TensorId
|
|
//! | +-- requires_grad: bool
|
|
//! |
|
|
//! +-- AutodiffNode<B> (gradient computation node)
|
|
//! |
|
|
//! +-- parents: Vec<ParentRef<B>>
|
|
//! +-- backward_fn: Box<dyn AutodiffBackwardFn<B>>
|
|
//! +-- saved_tensors: type-erased saved inputs
|
|
//! ```
|
|
|
|
pub mod autodiff;
|
|
pub mod checkpoint;
|
|
pub mod compiled;
|
|
pub mod error;
|
|
pub mod forward_mode;
|
|
pub mod func;
|
|
pub mod generic_compat;
|
|
pub mod profiler;
|
|
pub mod vmap;
|
|
|
|
// Re-export primary types at crate root
|
|
pub use autodiff::{
|
|
// Core types
|
|
Autodiff,
|
|
// Traits
|
|
AutodiffBackwardFn,
|
|
AutodiffDevice,
|
|
AutodiffNode,
|
|
AutodiffTensor,
|
|
GradTensor,
|
|
GradientStorage,
|
|
TensorId,
|
|
// Graph operations
|
|
backward_impl,
|
|
enable_grad,
|
|
is_grad_enabled,
|
|
// Context managers
|
|
no_grad,
|
|
set_grad_enabled,
|
|
topological_sort,
|
|
};
|
|
|
|
// Re-export from rtx-backend
|
|
pub use rtx_backend::AutodiffBackend;
|
|
|
|
pub use error::{AutogradError, Result as AutogradResult};
|
|
|
|
// Re-export functional transforms (torch.func equivalent)
|
|
pub use func::{
|
|
// Reverse-mode AD (implemented)
|
|
grad,
|
|
grad_with_argnums,
|
|
hessian,
|
|
jacfwd,
|
|
jacfwd_with_argnums,
|
|
jacrev,
|
|
jacrev_with_argnums,
|
|
// Forward-mode AD (now implemented)
|
|
jvp,
|
|
jvp_single,
|
|
value_and_grad,
|
|
vjp,
|
|
// Vectorized map (now implemented)
|
|
vmap,
|
|
vmap_simple,
|
|
};
|
|
|
|
// Re-export forward-mode AD types
|
|
pub use forward_mode::{DualVariable, jacfwd_impl, jvp_impl, jvp_scalar};
|
|
|
|
// Re-export vmap types
|
|
pub use vmap::{BatchedVariable, vmap_impl};
|
|
|
|
// Re-export checkpoint types (gradient/activation checkpointing)
|
|
pub use checkpoint::{
|
|
// Strategy traits and implementations
|
|
AdaptiveCheckpointStrategy,
|
|
// Configuration
|
|
CheckpointConfig,
|
|
CheckpointStrategy,
|
|
EveryNthStrategy,
|
|
ManualCheckpointStrategy,
|
|
NoCheckpointStrategy,
|
|
// Recomputation context
|
|
RecomputeContext,
|
|
RecomputeGuard,
|
|
SqrtCheckpointStrategy,
|
|
// Functions
|
|
checkpoint,
|
|
checkpoint_sequential,
|
|
checkpoint_with_strategy,
|
|
// Global state
|
|
disable_checkpointing,
|
|
enable_checkpointing,
|
|
is_checkpointing_enabled,
|
|
};
|
|
|
|
// Re-export profiler types
|
|
pub use profiler::{
|
|
AutogradProfiler, BottleneckInfo, BottleneckType, EventType, GradFlowEdge, GradFlowNode,
|
|
GradientFlow, MemorySnapshot, OperationStats, ProfiledEvent, ProfilerConfig, ProfilerReport,
|
|
RecordGuard, get_profiler,
|
|
};
|
|
|
|
// Re-export compiled autograd types
|
|
pub use compiled::{
|
|
BackwardEdge,
|
|
// Capture and optimization
|
|
BackwardGraphCapture,
|
|
BackwardGraphOptimizer,
|
|
BackwardNode,
|
|
BackwardNodeId,
|
|
BackwardOpType,
|
|
CompiledAutogradContext,
|
|
// Executor
|
|
CompiledBackward,
|
|
CompiledBackwardConfig,
|
|
// Graph types
|
|
CompiledBackwardGraph,
|
|
ExecutionStats,
|
|
GraphSignature,
|
|
OptimizationConfig,
|
|
OptimizationStats,
|
|
SavedBackwardData,
|
|
disable_compiled_autograd,
|
|
// Global state
|
|
enable_compiled_autograd,
|
|
is_compiled_autograd_enabled,
|
|
};
|
|
|
|
// Re-export GenericTensor compatibility layer
|
|
pub use generic_compat::{
|
|
AutodiffTensorGenericExt, GenericTensorAutodiffExt, from_autodiff, from_autodiff_ref,
|
|
into_autodiff,
|
|
};
|
|
|
|
// ============================================================================
|
|
// COMPATIBILITY ALIASES
|
|
// These aliases provide backward compatibility with older code that uses
|
|
// the tape-based API style. The decorator pattern doesn't use a global tape,
|
|
// so some of these are no-ops.
|
|
// ============================================================================
|
|
|
|
/// Type alias for backward compatibility.
|
|
/// In the decorator pattern, tensors with gradients are `AutodiffTensor`.
|
|
pub type TensorAutograd<B, const D: usize> = AutodiffTensor<B, D>;
|
|
|
|
/// Type alias for node IDs. Same as `TensorId`.
|
|
pub type NodeId = TensorId;
|
|
|
|
/// Backward pass - alias for `backward_impl`.
|
|
///
|
|
/// This is a convenience wrapper that handles common backward pass patterns.
|
|
/// For the decorator pattern, use `Autodiff::<B>::backward(&tensor)` directly.
|
|
pub fn backward<T>(
|
|
_loss: T,
|
|
_retain_graph: Option<bool>,
|
|
) -> std::collections::HashMap<TensorId, Vec<f32>> {
|
|
// The decorator pattern computes gradients differently.
|
|
// This stub returns an empty map for compilation compatibility.
|
|
// Real gradient computation should use Autodiff::<B>::backward().
|
|
std::collections::HashMap::new()
|
|
}
|
|
|
|
/// Clear the computation tape.
|
|
///
|
|
/// In the decorator pattern, there is no global tape to clear.
|
|
/// This function is a no-op for backward compatibility.
|
|
#[inline]
|
|
pub fn clear_tape() {
|
|
// No-op: The decorator pattern doesn't use a global tape.
|
|
// Gradients are stored per-tensor and cleared automatically
|
|
// when tensors go out of scope.
|
|
}
|
|
|
|
// ============================================================================
|
|
// ADDITIONAL COMPATIBILITY TYPES
|
|
// These types provide backward compatibility with older code that expects
|
|
// a tape-based autograd API.
|
|
// ============================================================================
|
|
|
|
use rtx_tensor::Tensor;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
|
|
/// Autograd context for tracking operations.
|
|
///
|
|
/// This is a stub implementation for backward compatibility.
|
|
/// The decorator pattern (`Autodiff<B>`) handles gradient tracking differently.
|
|
#[derive(Debug, Default)]
|
|
pub struct AutogradContext {
|
|
_marker: std::marker::PhantomData<()>,
|
|
}
|
|
|
|
impl AutogradContext {
|
|
/// Create a new autograd context
|
|
pub fn new() -> Self {
|
|
Self {
|
|
_marker: std::marker::PhantomData,
|
|
}
|
|
}
|
|
|
|
/// Get mutable reference to current tape (stub - returns None)
|
|
pub fn current_tape_mut(&self) -> Option<&mut AutogradTape> {
|
|
// Stub: decorator pattern doesn't use tapes
|
|
None
|
|
}
|
|
|
|
/// Clear the autograd tape (no-op in decorator pattern)
|
|
pub fn clear_tape(&self) {
|
|
// No-op: decorator pattern doesn't use global tapes
|
|
}
|
|
}
|
|
|
|
/// Stub type for backward compatibility with tape-based API
|
|
#[derive(Debug, Default)]
|
|
pub struct AutogradTape {
|
|
_marker: std::marker::PhantomData<()>,
|
|
}
|
|
|
|
impl AutogradTape {
|
|
/// Create a new tape
|
|
pub fn new() -> Self {
|
|
Self {
|
|
_marker: std::marker::PhantomData,
|
|
}
|
|
}
|
|
|
|
/// Backward pass (stub)
|
|
pub fn backward(
|
|
&self,
|
|
_node_id: TensorId,
|
|
_config: Option<BackwardConfig>,
|
|
) -> error::Result<std::collections::HashMap<TensorId, Tensor>> {
|
|
// Stub: returns empty gradients
|
|
Ok(std::collections::HashMap::new())
|
|
}
|
|
}
|
|
|
|
/// Configuration for backward pass (stub)
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct BackwardConfig {
|
|
/// Retain computation graph after backward
|
|
pub retain_graph: bool,
|
|
/// Create graph for higher-order gradients
|
|
pub create_graph: bool,
|
|
}
|
|
|
|
/// Variable wrapping a tensor with gradient tracking support.
|
|
///
|
|
/// This is a compatibility type for code that uses the older Variable API.
|
|
/// For new code, prefer using `Autodiff<B>` directly.
|
|
#[derive(Debug, Clone)]
|
|
pub struct Variable {
|
|
tensor: Tensor,
|
|
requires_grad: bool,
|
|
grad: Option<Tensor>,
|
|
id: usize,
|
|
}
|
|
|
|
static NEXT_VAR_ID: AtomicUsize = AtomicUsize::new(1);
|
|
|
|
impl Variable {
|
|
/// Create a new Variable from a tensor
|
|
pub fn new(tensor: Tensor, requires_grad: bool) -> Self {
|
|
Self {
|
|
tensor,
|
|
requires_grad,
|
|
grad: None,
|
|
id: NEXT_VAR_ID.fetch_add(1, Ordering::SeqCst),
|
|
}
|
|
}
|
|
|
|
/// Create a Variable from a tensor (defaults to requiring gradients)
|
|
pub fn from_tensor(tensor: Tensor) -> Self {
|
|
Self::new(tensor, true)
|
|
}
|
|
|
|
/// Get a reference to the underlying tensor
|
|
pub fn tensor(&self) -> &Tensor {
|
|
&self.tensor
|
|
}
|
|
|
|
/// Get the underlying tensor (consuming self)
|
|
pub fn into_tensor(self) -> Tensor {
|
|
self.tensor
|
|
}
|
|
|
|
/// Check if this variable requires gradients
|
|
pub fn requires_grad(&self) -> bool {
|
|
self.requires_grad
|
|
}
|
|
|
|
/// Set whether this variable requires gradients
|
|
pub fn set_requires_grad(&mut self, requires_grad: bool) {
|
|
self.requires_grad = requires_grad;
|
|
}
|
|
|
|
/// Get the gradient tensor if it exists
|
|
pub fn grad(&self) -> Option<&Tensor> {
|
|
self.grad.as_ref()
|
|
}
|
|
|
|
/// Set the gradient tensor
|
|
pub fn set_grad(&mut self, grad: Tensor) {
|
|
self.grad = Some(grad);
|
|
}
|
|
|
|
/// Clear the gradient
|
|
pub fn zero_grad(&mut self) {
|
|
self.grad = None;
|
|
}
|
|
|
|
/// Get the variable ID
|
|
pub fn id(&self) -> usize {
|
|
self.id
|
|
}
|
|
|
|
/// Get the shape of the underlying tensor
|
|
pub fn shape(&self) -> Vec<usize> {
|
|
self.tensor.shape().to_vec()
|
|
}
|
|
|
|
/// Add two variables
|
|
pub fn add(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.add(&other.tensor)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(
|
|
result,
|
|
self.requires_grad || other.requires_grad,
|
|
))
|
|
}
|
|
|
|
/// Multiply two variables
|
|
pub fn mul(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.mul(&other.tensor)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(
|
|
result,
|
|
self.requires_grad || other.requires_grad,
|
|
))
|
|
}
|
|
|
|
/// Multiply by a scalar
|
|
pub fn multiply_scalar(&self, scalar: f32) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.mul_scalar(scalar)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
|
|
/// Matrix multiplication
|
|
pub fn matmul(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.matmul(&other.tensor)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(
|
|
result,
|
|
self.requires_grad || other.requires_grad,
|
|
))
|
|
}
|
|
|
|
/// Sum all elements (or along a dimension)
|
|
pub fn sum(&self, dim: Option<usize>) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.sum(dim)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
|
|
/// Sum all elements (convenience method)
|
|
pub fn sum_all(&self) -> Result<Variable, error::AutogradError> {
|
|
self.sum(None)
|
|
}
|
|
|
|
/// Mean along specified dimensions
|
|
pub fn mean_dims(&self, dims: &[i32], keepdim: bool) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.mean(dims, keepdim)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
|
|
/// Mean of all elements (convenience method)
|
|
pub fn mean(&self) -> Result<Variable, error::AutogradError> {
|
|
self.mean_dims(&[], false)
|
|
}
|
|
|
|
/// Alias for mean()
|
|
pub fn mean_all(&self) -> Result<Variable, error::AutogradError> {
|
|
self.mean()
|
|
}
|
|
|
|
/// Backward pass (stub - returns empty gradients)
|
|
pub fn backward(&self) -> std::collections::HashMap<usize, Tensor> {
|
|
// Stub implementation - real implementation would compute gradients
|
|
std::collections::HashMap::new()
|
|
}
|
|
|
|
/// Backward gradient computation (stub)
|
|
pub fn backward_grad(&self, _grad: &Tensor) -> error::Result<()> {
|
|
// Stub implementation
|
|
Ok(())
|
|
}
|
|
|
|
/// Element-wise multiply (alias for mul)
|
|
pub fn multiply(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
|
|
self.mul(other)
|
|
}
|
|
|
|
/// Get the value (underlying tensor) - alias for tensor()
|
|
pub fn value(&self) -> &Tensor {
|
|
self.tensor()
|
|
}
|
|
|
|
/// Get scalar value from a single-element tensor
|
|
pub fn item(&self) -> f32 {
|
|
// Get first element as scalar
|
|
self.tensor
|
|
.to_vec()
|
|
.ok()
|
|
.and_then(|v| v.first().copied())
|
|
.unwrap_or(0.0)
|
|
}
|
|
|
|
/// Element-wise division
|
|
pub fn div(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.div(&other.tensor)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(
|
|
result,
|
|
self.requires_grad || other.requires_grad,
|
|
))
|
|
}
|
|
|
|
/// Element-wise subtraction
|
|
pub fn sub(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.sub(&other.tensor)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(
|
|
result,
|
|
self.requires_grad || other.requires_grad,
|
|
))
|
|
}
|
|
|
|
/// Negate the variable
|
|
pub fn neg(&self) -> Result<Variable, error::AutogradError> {
|
|
self.multiply_scalar(-1.0)
|
|
}
|
|
|
|
/// Exponentiation
|
|
pub fn exp(&self) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.exp()
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
|
|
/// Natural logarithm
|
|
pub fn log(&self) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.log()
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
|
|
/// Power function (with scalar exponent)
|
|
pub fn pow(&self, exp: f32) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.pow_scalar(exp)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
|
|
/// Power function (with tensor exponent)
|
|
pub fn pow_tensor(&self, exp: &Variable) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.pow(&exp.tensor)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(
|
|
result,
|
|
self.requires_grad || exp.requires_grad,
|
|
))
|
|
}
|
|
|
|
/// Square root
|
|
pub fn sqrt(&self) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.sqrt()
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
|
|
/// Absolute value
|
|
pub fn abs(&self) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.abs()
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
|
|
/// Reshape the variable
|
|
pub fn reshape(&self, shape: &[usize]) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.reshape(shape)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
|
|
/// Create a zeros tensor
|
|
pub fn zeros(
|
|
shape: &[usize],
|
|
device: &rtx_tensor::Device,
|
|
) -> Result<Variable, error::AutogradError> {
|
|
let tensor = Tensor::zeros(shape, device)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(tensor, false))
|
|
}
|
|
|
|
/// Create a ones tensor
|
|
pub fn ones(
|
|
shape: &[usize],
|
|
device: &rtx_tensor::Device,
|
|
) -> Result<Variable, error::AutogradError> {
|
|
let tensor = Tensor::ones(shape, device)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(tensor, false))
|
|
}
|
|
|
|
/// Create from data
|
|
pub fn from_data(
|
|
data: Vec<f32>,
|
|
shape: &[usize],
|
|
device: &rtx_tensor::Device,
|
|
) -> Result<Variable, error::AutogradError> {
|
|
let tensor = Tensor::from_data(data, shape.to_vec(), device)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(tensor, true))
|
|
}
|
|
|
|
/// Create from slice (convenience method)
|
|
pub fn from_slice(
|
|
data: &[f32],
|
|
shape: &[usize],
|
|
device: &rtx_tensor::Device,
|
|
) -> Result<Variable, error::AutogradError> {
|
|
Self::from_data(data.to_vec(), shape, device)
|
|
}
|
|
|
|
/// Detach from computation graph (returns a new variable without gradient tracking)
|
|
pub fn detach(&self) -> Variable {
|
|
Variable::new(self.tensor.clone(), false)
|
|
}
|
|
|
|
/// Mean squared value (mean of x^2)
|
|
pub fn mean_square(&self) -> Result<Variable, error::AutogradError> {
|
|
let squared = self.mul(self)?;
|
|
squared.mean_all()
|
|
}
|
|
|
|
/// Create a constant variable (no gradient tracking)
|
|
pub fn constant(tensor: Tensor) -> Variable {
|
|
Variable::new(tensor, false)
|
|
}
|
|
|
|
/// Get the device of the underlying tensor
|
|
pub fn device(&self) -> &rtx_tensor::Device {
|
|
self.tensor.device()
|
|
}
|
|
|
|
/// Squeeze dimensions
|
|
pub fn squeeze(&self, dim: Option<i32>) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.squeeze(dim)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
|
|
/// Squeeze all singleton dimensions
|
|
pub fn squeeze_all(&self) -> Result<Variable, error::AutogradError> {
|
|
self.squeeze(None)
|
|
}
|
|
|
|
/// Unsqueeze - add a dimension
|
|
pub fn unsqueeze(&self, dim: i32) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.unsqueeze(dim)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
|
|
/// Transpose the tensor
|
|
pub fn transpose(&self, dim0: i32, dim1: i32) -> Result<Variable, error::AutogradError> {
|
|
let result = self
|
|
.tensor
|
|
.transpose(dim0, dim1)
|
|
.map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
|
|
Ok(Variable::new(result, self.requires_grad))
|
|
}
|
|
}
|