tinygrad/
lib.rs

1//! # tinygrad
2//!
3//! `tinygrad` is a crate for building and training neural networks in Rust. It provides a simple interface for defining tensors,
4//! performing forward and backward passes, and implementing basic operations such as dot products and summation.
5//!
6//! ## Quick Start
7//!
8//! Get started with the `tinygrad` library by following these simple steps:
9//!
10//! 1. Install the `tinygrad` crate by adding the following line to your `Cargo.toml` file:
11//!
12//! ```toml
13//! [dependencies]
14//! tinygrad = "0.1.0"
15//! ```
16//!
17//! 2. Use the `Tensor` and `ForwardBackward` traits to create and work with tensors:
18//!
19//! ```rust
20//! use ndarray::{array, Array1};
21//! use tinygrad::{Tensor, Context, TensorTrait};
22//!
23//! // Create a tensor
24//! let value = array![1.0, 2.0, 3.0];
25//! let tensor = Tensor::new(value);
26//!
27//! // Perform forward and backward passes
28//! let mut ctx = Context::new();
29//! let result = tensor.forward(&mut ctx, vec![tensor.get_value()]);
30//! tensor.backward(&mut ctx, array![1.0, 1.0, 1.0].view());
31//! ```
32//!
33//! 3. Implement custom operations by defining structs that implement the `ForwardBackward` trait:
34//!
35//! ```rust
36//! use ndarray::ArrayView1;
37//! use tinygrad::{ForwardBackward, Context, TensorTrait};
38//!
39//! // Example operation: Dot product
40//! struct Dot;
41//!
42//! impl ForwardBackward for Dot {
43//!     fn forward(&self, _ctx: &mut Context, inputs: Vec<ArrayView1<f64>>) -> f64 {
44//!         let input = &inputs[0];
45//!         let weight = &inputs[1];
46//!         input.dot(weight)
47//!     }
48//!
49//!     fn backward(&self, ctx: &mut Context, grad_output: ArrayView1<f64>) {
50//!         // Implement backward pass
51//!         // ...
52//!     }
53//! }
54//! ```
55//!
56//! ## GitHub Repository
57//!
58//! You can access the source code for the `tinygrad` crate on [GitHub](https://github.com/wiseaidev/tinygrad).
59//!
60//! ## Contributing
61//!
62//! Contributions and feedback are welcome! If you'd like to contribute, report an issue, or suggest an enhancement,
63//! please engage with the project on [GitHub](https://github.com/wiseaidev/tinygrad).
64//! Your contributions help improve this crate for the community.
65
66use ndarray::{array, Array1, ArrayView1};
67
68type Gradient = Array1<f64>;
69
70/// This trait defines the common interface for tensors in a computational graph.
71pub trait TensorTrait {
72    /// Computes the forward pass of the tensor.
73    ///
74    /// # Arguments
75    /// * `ctx` - A mutable reference to the computation context.
76    /// * `inputs` - A vector of input arrays for the forward pass.
77    ///
78    /// # Returns
79    /// (`f64`): The result of the forward pass.
80    fn forward(&self, ctx: &mut Context, inputs: Vec<ArrayView1<f64>>) -> f64;
81
82    /// Computes the backward pass of the tensor to calculate gradients.
83    ///
84    /// # Arguments
85    /// * `ctx` - A mutable reference to the computation context.
86    /// * `grad_output` - The gradient of the loss with respect to the output.
87    fn backward(&self, ctx: &mut Context, grad_output: ArrayView1<f64>);
88
89    /// Gets the value of the tensor.
90    ///
91    /// # Returns
92    /// (`ArrayView1<f64>`): The view of the tensor's value.
93    fn get_value(&self) -> ArrayView1<f64>;
94
95    /// Gets the gradient of the tensor.
96    ///
97    /// # Returns
98    /// (`Option<Gradient>`): The option containing the gradient if available.
99    fn get_grad(&self) -> Option<Gradient>;
100}
101
102/// Represents a basic implementation of a tensor.
103#[derive(Debug, Clone)]
104pub struct Tensor {
105    pub value: Array1<f64>,
106    pub grad: Option<Gradient>,
107}
108
109impl Tensor {
110    /// Creates a new Tensor with the given value.
111    ///
112    /// # Arguments
113    ///
114    /// * `value` - The array representing the value of the tensor.
115    ///
116    /// # Returns
117    ///
118    /// (`Tensor`): A new Tensor instance.
119    ///
120    /// # Examples
121    ///
122    /// ```
123    /// use ndarray::{array, Array1};
124    /// use tinygrad::Tensor;
125    ///
126    /// let value = array![1.0, 2.0, 3.0];
127    /// let tensor = Tensor::new(value);
128    /// ```
129    pub fn new(value: Array1<f64>) -> Tensor {
130        Tensor { value, grad: None }
131    }
132}
133
134impl TensorTrait for Tensor {
135    /// Implements the forward pass for the Tensor.
136    ///
137    /// # Arguments
138    ///
139    /// * `_ctx` - A mutable reference to the computation context (not used in this example).
140    /// * `_inputs` - A vector of ArrayView1<f64> representing the input values (not used in this example).
141    ///
142    /// # Returns
143    ///
144    /// (`f64`): The result of the forward pass (not implemented in this example, returns 0.0).
145    fn forward(&self, _ctx: &mut Context, _inputs: Vec<ArrayView1<f64>>) -> f64 {
146        // TODO: Implement forward function for Tensor
147        0.0
148    }
149
150    /// Implements the backward pass for the Tensor.
151    ///
152    /// # Arguments
153    ///
154    /// * `_ctx` - A mutable reference to the computation context (not used in this example).
155    /// * `_grad_output` - An ArrayView1<f64> representing the gradient of the output.
156    ///
157    /// # Examples
158    ///
159    /// ```
160    /// use ndarray::{array, Array1};
161    /// use tinygrad::{Tensor, Context, TensorTrait};
162    ///
163    /// let mut ctx = Context::new();
164    /// let tensor = Tensor::new(array![1.0, 2.0, 3.0]);
165    /// tensor.backward(&mut ctx, array![1.0, 1.0, 1.0].view());
166    /// ```
167    fn backward(&self, _ctx: &mut Context, _grad_output: ArrayView1<f64>) {
168        // TODO: Implement backward function for Tensor
169    }
170
171    /// Returns the value of the Tensor.
172    ///
173    /// # Returns
174    ///
175    /// (`ArrayView1<f64>`): The view of the array representing the value of the tensor.
176    ///
177    /// # Examples
178    ///
179    /// ```
180    /// use ndarray::{array, Array1};
181    /// use tinygrad::{Tensor, TensorTrait};
182    ///
183    /// let tensor = Tensor::new(array![1.0, 2.0, 3.0]);
184    /// let value = tensor.get_value();
185    /// ```
186    fn get_value(&self) -> ArrayView1<f64> {
187        self.value.view()
188    }
189
190    /// Returns the gradient of the Tensor if available.
191    ///
192    /// # Returns
193    ///
194    /// (`Option<Gradient>`): The optional gradient of the tensor.
195    ///
196    /// # Examples
197    ///
198    /// ```
199    /// use ndarray::{array, Array1};
200    /// use tinygrad::{Tensor, TensorTrait};
201    ///
202    /// let tensor = Tensor::new(array![1.0, 2.0, 3.0]);
203    /// let grad = tensor.get_grad();
204    /// ```
205    fn get_grad(&self) -> Option<Gradient> {
206        self.grad.clone()
207    }
208}
209
210/// This trait defines the interface for operations that have both forward and backward passes.
211pub trait ForwardBackward {
212    /// Computes the forward pass of the operation.
213    ///
214    /// # Arguments
215    /// * `ctx` - A mutable reference to the computation context.
216    /// * `inputs` - A vector of input arrays for the forward pass.
217    ///
218    /// # Returns
219    /// (`f64`): The result of the forward pass.
220    fn forward(&self, ctx: &mut Context, inputs: Vec<ArrayView1<f64>>) -> f64;
221
222    /// Computes the backward pass of the operation to calculate gradients.
223    ///
224    /// # Arguments
225    /// * `ctx` - A mutable reference to the computation context.
226    /// * `grad_output` - The gradient of the loss with respect to the output.
227    fn backward(&self, ctx: &mut Context, grad_output: ArrayView1<f64>);
228}
229
230/// Represents the Dot operation.
231struct Dot;
232
233/// Represents the Sum operation.
234struct Sum;
235
236impl ForwardBackward for Dot {
237    fn forward(&self, _ctx: &mut Context, inputs: Vec<ArrayView1<f64>>) -> f64 {
238        let input = &inputs[0];
239        let weight = &inputs[1];
240
241        input.dot(weight)
242    }
243
244    fn backward(&self, ctx: &mut Context, grad_output: ArrayView1<f64>) {
245        if ctx.saved_tensors.is_empty() {
246            println!("Warning: saved_tensors is empty. Unable to compute gradients.");
247            return;
248        }
249
250        let mut input = ctx.saved_tensors[0].clone();
251        let mut weight = ctx.saved_tensors[1].clone();
252
253        let grad_input = grad_output.dot(&input.get_value().t());
254        let grad_weight = input.get_value().t().dot(&grad_output);
255
256        input.grad = Some(array![grad_input]);
257        weight.grad = Some(array![grad_weight]);
258
259        ctx.save_for_backward(vec![Box::new(*input), Box::new(*weight)]);
260    }
261}
262
263impl ForwardBackward for Sum {
264    fn forward(&self, _ctx: &mut Context, inputs: Vec<ArrayView1<f64>>) -> f64 {
265        let input = &inputs[0];
266
267        input.sum()
268    }
269
270    fn backward(&self, ctx: &mut Context, grad_output: ArrayView1<f64>) {
271        let mut input = ctx.saved_tensors[0].clone();
272
273        input.grad = Some(Array1::from(grad_output.map(|x| x * 1.0)));
274
275        ctx.save_for_backward(vec![Box::new(*input)]);
276    }
277}
278
279/// Represents the computation context, storing tensors for backward pass computations.
280pub struct Context {
281    pub saved_tensors: Vec<Box<Tensor>>,
282}
283
284impl Context {
285    /// Creates a new Context instance.
286    ///
287    /// # Returns
288    ///
289    /// (`Context`): A new Context instance.
290    ///
291    /// # Examples
292    ///
293    /// ```
294    /// use tinygrad::Context;
295    ///
296    /// let context = Context::new();
297    /// ```
298    pub fn new() -> Context {
299        Context {
300            saved_tensors: Vec::new(),
301        }
302    }
303
304    /// Saves tensors for backward pass.
305    ///
306    /// # Arguments
307    ///
308    /// * `tensors` - A vector of Boxed Tensors to be saved.
309    ///
310    /// # Examples
311    ///
312    /// ```
313    /// use ndarray::{array, Array1};
314    /// use tinygrad::{Tensor, Context};
315    ///
316    /// let mut ctx = Context::new();
317    /// let tensor = Tensor::new(array![1.0, 2.0, 3.0]);
318    /// ctx.save_for_backward(vec![Box::new(tensor)]);
319    /// ```
320    pub fn save_for_backward(&mut self, tensors: Vec<Box<Tensor>>) {
321        self.saved_tensors.extend(tensors);
322    }
323}
324
325impl Default for Context {
326    fn default() -> Self {
327        Self::new()
328    }
329}
tinygrad/lib.rs

tinygrad/
lib.rs
