Add some documentation and test to the linear layer. (#151)

* Add some documentation and test to the linear layer.

* Layer norm doc.

* Minor tweaks.

README.md

@@ -31,7 +31,9 @@ Cheatsheet:
 
 |            | Using PyTorch                            | Using Candle                                                     |
 |------------|------------------------------------------|------------------------------------------------------------------|
-| Creation   | `torch.Tensor([[1, 2], [3, 4]])`         | `Tensor::new(&[[1f32, 2.]], [3., 4.]], &Device::Cpu)?`           |
+| Creation   | `torch.Tensor([[1, 2], [3, 4]])`         | `Tensor::new(`                                                   |
+|            |                                          | `  &[[1f32, 2.]], [3., 4.]],`                                    |
+|            |                                          | `  &Device::Cpu)?`                                               |
 | Indexing   | `tensor[:, :4]`                          | `tensor.i((.., ..4))?`                                           |
 | Operations | `tensor.view((2, 2))`                    | `tensor.reshape((2, 2))?`                                        |
 | Operations | `a.matmul(b)`                            | `a.matmul(&b)?`                                                  |

candle-nn/src/conv.rs

@@ -1,3 +1,4 @@
+//! Convolution Layers.
 use candle::{Result, Tensor};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]

candle-nn/src/embedding.rs

@@ -1,3 +1,4 @@
+//! Embedding Layer.
 use candle::{Result, Tensor};
 
 #[derive(Debug)]

candle-nn/src/layer_norm.rs

@@ -1,3 +1,33 @@
+//! Layer Normalization.
+//!
+//! This layer applies Layer Normalization over a mini-batch of inputs as described in [`Layer
+//! Normalization`]. The input is expected to have three dimensions: a batch dimension, a length,
+//! and a hidden size, the normalization is applied over the last dimension.
+//!
+//! # Example
+//!
+//! ```rust
+//! use candle::{Tensor, Device::Cpu};
+//! use candle_nn::LayerNorm;
+//! # fn main() -> candle::Result<()> {
+//!
+//! let w = Tensor::new(1f32, &Cpu)?;
+//! let b = Tensor::new(0f32, &Cpu)?;
+//! let layer = LayerNorm::new(w, b, 1e-5);
+//!
+//! let xs = Tensor::new(
+//!     &[[[1f32, 2., 3.], [4., 5., 6.], [9., 8., 7.]]],
+//!     &Cpu)?;
+//! let ys = layer.forward(&xs)?;
+//! assert_eq!(
+//!     ys.to_vec3::<f32>()?,
+//!     &[[[-1.2247356, 0.0,  1.2247356],
+//!        [-1.2247356, 0.0,  1.2247356],
+//!        [ 1.2247356, 0.0, -1.2247356]]]);
+//! # Ok(()) }
+//! ```
+//!
+//! [`Layer Normalization`]: https://arxiv.org/abs/1607.06450
 use candle::{DType, Result, Tensor};
 
 // This layer norm version handles both weight and bias so removes the mean.

candle-nn/src/linear.rs

@@ -1,3 +1,22 @@
+//! Linear layer
+//!
+//! This layer applies a linear transformation to the incoming data, `y = x@w.t() + b`.
+//! The bias is optional. The `forward` method can be used to apply the layer, it supports input
+//! with a batch dimension (so of shape `(b_sz, in_c)`) or without (of shape `(in_c,)`), the
+//! output has shape `(b_sz, out_c)` and `(out_c,)` respectively.
+//!
+//! ```rust
+//! use candle::{Tensor, Device::Cpu};
+//! use candle_nn::Linear;
+//! # fn main() -> candle::Result<()> {
+//!
+//! let w = Tensor::new(&[[1f32, 2.], [3., 4.], [5., 6.]], &Cpu)?;
+//! let layer = Linear::new(w, None); // Use no bias.
+//! let xs = Tensor::new(&[[10f32, 100.]], &Cpu)?;
+//! let ys = layer.forward(&xs)?;
+//! assert_eq!(ys.to_vec2::<f32>()?, &[[210.0, 430.0, 650.0]]);
+//! # Ok(()) }
+//! ```
 use candle::Tensor;
 
 #[derive(Debug)]