hasktorch-CudaHashMap-0.1.0.0: test/Main.hs

module Main (main) where

import Data.Int (Int32, Int64)
import Torch.Functional (indexPut)
import Torch.Tensor (asTensor, device, indexSelect, toDevice)
import Torch.TensorFactories (zeros')
import Torch.CudaHashMap

main :: IO ()
main = do
  putStrLn "=== Test 1: Int32 Key -> Int32 Value ==="
  m1 <- newCudaHashMap
  let k1 = asTensor ([1, 2, 3] :: [Int32])
      v1 = asTensor ([10, 20, 30] :: [Int32])
      v1_overwrite = asTensor ([100, 200, 300] :: [Int32])
  insertBatch m1 k1 v1
  insertBatch m1 k1 v1_overwrite
  finalizeCudaHashMap m1
  res1 <- lookupBatch m1 (asTensor ([3, 1, 99] :: [Int32]))
  print res1

  putStrLn "=== Test 2: Int64 Key -> Float Value ==="
  m2 <- newCudaHashMap
  let k2 = asTensor ([100, 200, 300] :: [Int64])
      v2 = asTensor ([1.25, 2.5, 3.75] :: [Float])
  insertBatch m2 k2 v2
  finalizeCudaHashMap m2
  res2 <- lookupBatch m2 (asTensor ([200, 100, 999] :: [Int64]))
  print res2

  putStrLn "=== Test 3: Key -> Index (Int64) for indexPut / indexSelect ==="
  m3 <- newCudaHashMap
  let table0 = zeros' [10]
      batchKeys1 = asTensor ([500, 600, 500, 700] :: [Int64])
      batchVals1 = asTensor ([5.0, 6.0, 5.5, 7.0] :: [Float])
  idx1 <- insertKeys m3 batchKeys1
  let dev    = device idx1
      table1 = indexPut False [idx1] (toDevice dev batchVals1) (toDevice dev table0)
  finalizeCudaHashMap m3
  qIdx <- lookupIndices m3 (asTensor ([700, 500, 600] :: [Int64]))
  let out = indexSelect 0 qIdx table1
  print out