Lenses

Hasktorch uses lenses in two directions: inward, addressing parts of a single tensor (slices, named fields), and outward, traversing every tensor inside an arbitrary structure. Both come in untyped and typed flavors. This chapter walks through all four quadrants with executable examples.

import Torch
import Torch.Index (lslice, slice)
import Torch.Lens (types, flattenValues, replaceValues)
import Lens.Family
import GHC.Generics (Generic)
import qualified Torch.Typed as T
import qualified Torch.Typed.Lens as TL
import qualified Torch.Typed.Index as TI
import Data.Vector.Sized (Vector)

Slicing lenses (untyped)

The lslice quasiquoter turns PyTorch indexing syntax into a lens into a tensor: something you can read through, but also write through.

let t = asTensor ([[0, 1, 2], [3, 4, 5]] :: [[Float]])

Reading a slice with ^.:

t ^. [lslice|1|] :: Tensor => 
Tensor Float [3] [ 3.0000   ,  4.0000   ,  5.0000   ]

t ^. [lslice|:, ::2|] :: Tensor => 
Tensor Float [2,2] [[ 0.0000,  2.0000   ],
                    [ 3.0000   ,  5.0000   ]]

Writing through the same lens with .~ returns an updated tensor, with everything outside the slice untouched:

t & [lslice|0|] .~ zeros' [3] => 
Tensor Float [2,3] [[ 0.0000,  0.0000,  0.0000],
                    [ 3.0000   ,  4.0000   ,  5.0000   ]]

And %~ modifies the focused slice with a function:

t & [lslice|:, 1|] %~ ((* 100) :: Tensor -> Tensor) => 
Tensor Float [2,3] [[ 0.0000,  100.0000   ,  2.0000   ],
                    [ 3.0000   ,  400.0000   ,  5.0000   ]]

This is the lens counterpart of t[0] = ... in Python — but pure: the original t is unchanged.

Structure traversals (untyped)

Torch.Lens provides types, a generic traversal of every value of a given type inside a structure. Any record with a Generic instance works.

data TwoLayer = TwoLayer
  { weight1 :: Tensor
  , weight2 :: Tensor
  , steps   :: Int
  } deriving (Generic, Show)

let m = TwoLayer (ones' [2, 3]) (ones' [3]) 0

Collect every tensor:

map shape (flattenValues (types @Tensor) m) => 
[[2,3],[3]]

Rewrite every tensor — this is exactly how toType and toDevice convert whole models; here we switch the model to half precision:

map dtype (flattenValues (types @Tensor) (toType Half m)) => 
[Half,Half]

Note the Int field rides along untouched; only the tensors are visited.

Field and shape lenses (typed)

In the typed API, dimensions can be records, and record fields become lenses. field @"r" reads or writes the r component of an RGB dimension, and its very existence is checked at compile time — field @"q" would not compile.

data RGB a = RGB { r :: a, g :: a, b :: a } deriving (Generic, Show)

let img = T.fromUnnamed T.ones :: T.NamedTensor '( 'CPU, 0) 'Float '[Vector 2, RGB]

T.toDynamic (img ^. TL.field @"r") => 
Tensor Float [2] [ 1.0000   ,  1.0000   ]

T.toDynamic (img & TL.field @"g" .~ T.fromUnnamed T.zeros) => 
Tensor Float [2,3] [[ 1.0000   ,  0.0000,  1.0000   ],
                    [ 1.0000   ,  0.0000,  1.0000   ]]

Slicing lenses exist here too: sliceLens is the typed lslice, built from getSlice and setSlice, and the same slice quasiquoter works in type position. The focused shape is computed at compile time, so writing a wrongly-shaped value through the lens does not compile.

let tt = T.ones :: T.Tensor '( 'CPU, 0) 'Float '[2, 3]

T.toDynamic (tt ^. TI.sliceLens @[slice| 1 |]) => 
Tensor Float [3] [ 1.0000   ,  1.0000   ,  1.0000   ]

T.toDynamic (tt & TI.sliceLens @[slice| :, 1 |] %~ (* 100)) => 
Tensor Float [2,3] [[ 1.0000   ,  100.0000   ,  1.0000   ],
                    [ 1.0000   ,  100.0000   ,  1.0000   ]]

The outward traversal also exists in typed form, where it becomes shape-selective: types at a typed tensor target visits only the tensors of exactly that device, dtype and shape.

data TypedNet device = TypedNet
  { l1 :: T.Tensor device 'Float '[2, 3]
  , l2 :: T.Tensor device 'Float '[3, 4]
  , l3 :: T.Tensor device 'Float '[2, 3]
  } deriving Generic

let net = TypedNet T.ones T.ones T.zeros :: TypedNet '( 'CPU, 0)

Only the two '[2, 3] layers are visited; the '[3, 4] one is not:

length (flattenValues (types @(T.Tensor '( 'CPU, 0) 'Float '[2, 3])) net) => 
2

T.toDynamic (l3 (over (types @(T.Tensor '( 'CPU, 0) 'Float '[2, 3])) (+ 1) net)) => 
Tensor Float [2,3] [[ 1.0000   ,  1.0000   ,  1.0000   ],
                    [ 1.0000   ,  1.0000   ,  1.0000   ]]

replaceValues swaps specific structures wholesale — replace the two '[2, 3] tensors and leave everything else:

T.toDynamic (l1 (replaceValues (types @(T.Tensor '( 'CPU, 0) 'Float '[2, 3])) net [T.zeros, T.ones])) => 
Tensor Float [2,3] [[ 0.0000,  0.0000,  0.0000],
                    [ 0.0000,  0.0000,  0.0000]]

Retyping conversions (typed)

Finally, the typed counterparts of toType/toDevice change the type along with the values. Torch.Typed.DType.toDType converts every tensor in a model and rewrites the model's dtype parameter, so after converting to 'Half a stray 'Float batch is a compile-time error rather than a silent upcast:

model                                :: Linear 10 1 'Float '( 'CPU, 0)
toDType @'Half @'Float model         :: Linear 10 1 'Half  '( 'CPU, 0)
toDevice @'( 'CUDA, 0) @'( 'CPU, 0) model
                                     :: Linear 10 1 'Float '( 'CUDA, 0)

See Named Tensors and Lenses for the reference treatment of these conversions.