Module dfdx::tensor

source ·
Expand description

The Tensor struct, Cpu & Cuda devices, and traits like ZerosTensor, OnesTensor, SampleTensor.

At a high level a tensor is made up of:

  1. The crate::shapes::Shape of the array it stores
  2. The crate::shapes::Dtype of the elements of the array
  3. The DeviceStorage (e.g. Cpu or Cuda) that it uses to store the nd array
  4. A Tape, which can either actually be a tape (OwnedTape) or be empty (NoneTape).

Which are all generic parameters of Tensor. See the type’s docstring for more info

Creating a device

In order to do anything with tensors, you first need to construct the device that they will be stored on:

let dev: Cpu = Default::default();
let dev: Cpu = Cpu::seed_from_u64(0);
ⓘ
let dev: Cuda = Default::default();
let dev: Cuda = Cuda::seed_from_u64(1234);
let dev: Cuda = Cuda::try_build(0, 1234).unwrap();

Creating tensors

From arrays/vecs

See TensorFrom & TensorFromVec.

let _ = dev.tensor([1.0, 2.0, 3.0]);
let _ = dev.tensor_from_vec(vec![1.0, 2.0, 3.0], (3, ));

Filled with 0s or 1s

See ZerosTensor and OnesTensor.

let _: Tensor<Rank1<5>,f32 , _> = dev.zeros();
let _: Tensor<Rank2<3, 2>, f32, _> = dev.ones();

Filled with random data

See SampleTensor

let _: Tensor<Rank1<5>, f32, _> = dev.sample_uniform();
let _: Tensor<Rank2<3, 5>, f32, _> = dev.sample_normal();
// or pass in actual distributions
let _: Tensor<Rank1<3>, f32, _> = dev.sample(rand_distr::Standard);
let _: Tensor<Rank2<4, 3>, f32, _> = dev.sample(rand_distr::StandardNormal);

Copy data from slices

You can use Tensor::copy_from and Tensor::copy_into to copy data into a tensor:

let mut a: Tensor<Rank1<1000>, f32, _> = dev.zeros();
let buf: Vec<f32> = vec![1.0; 1000];
a.copy_from(&buf);

Modifying an already constructed tensor

There are only a few ways to do this, as normally you should just create a new tensor with tensor_ops.

See Tensor::fill_with_zeros, Tensor::fill_with_ones, Tensor::fill_with_distr

Converting tensors to rust arrays

Since the way tensors are stored is opaque to users (driven by whatever device the tensor is stored on), use the AsArray trait to convert tensors to actual rust arrays if you want to work with them directly.

let t: Tensor<Rank2<2, 3>, f32, _> = dev.zeros();
let t: [[f32; 3]; 2] = t.array();

Tracing gradients

Use the Tensor::trace or Tensor::traced methods to add OwnedTape to the Tensor. .trace() will clone the data, while .traced() will take ownership of the tensor and return a version with an OwnedTape.

Note that these two methods are only present for tensors without a tape already.

let t: Tensor<Rank1<5>,f32, _> = dev.zeros();
let mut grads = t.alloc_grads();
let t_clone: Tensor<Rank1<5>, f32, _, OwnedTape<f32, Cpu>> = t.trace(grads);

Gradient Accumulation

If you re-use the same gradients object without zero-ing out the gradients, you can implement gradient accumulation!

Serialization using numpy

See Tensor::save_to_npy and Tensor::load_from_npy.

You can also use Tensor::write_to_npz and Tensor::read_from_npz when working with zip archives.

Modules

Structs

  • Cpu
    A device that stores data on the heap.
  • Cuda
    A Cuda device that enables constructing tensors on GPUs & running GPU kernels.
  • Gradients
    A generic container for keeping gradients of tensors keyed by the tensor’s UniqueId.
  • NoneTape
    Contains nothing. When Tape::add_backward_op is called, this struct does nothing.
  • OwnedTape
    Contains a Gradients and list of backward operations.
  • Tensor
    The single tensor struct that stores nd arrays and tapes.
  • UniqueId
    An id used in to associate gradients with Tensors.

Enums

Traits

  • AsArray
  • CopySlice
    Enables copying data into and out of tensors
  • DeviceStorage
    Something that can store nd arrays for a given Shape and Dtype
  • HasErr
    Represents something that has an error associated type
  • Merge
    Combine two things
  • OnesTensor
    Construct tensors filled with ones.
  • PutTape
    Put a tape of type T into the tensor
  • SampleTensor
    Constructs tensors filled with random values from a given distribution.
  • SplitTape
    Remove the tape from a tensor
  • Tape
    Something that can track backward operations.
  • TensorFrom
    Construct tensors from rust data
  • TensorFromVec
    Construct tensors from rust vectors. This trait is only used to implement TensorFrom.
  • Trace
    Something that can trace gradients
  • WithEmptyTape
    Clones self and inserts a new empty tape into the clone
  • ZerosTensor
    Construct tensors filled with zeros.

Type Definitions