search

ctongfei / nexus

Experimental tensor-typed deep learning
https://tongfei.me/nexus/
MIT License
257 stars 15 forks source link

Nd4j backend #14

Closed drewnoff closed 5 years ago

drewnoff commented 6 years ago

An attempt to bring GPU support via nd4j backend.

One can switch between nd4j backends by choosing appropriate in sbt file:

lazy val nd4jBackend = Libs.nd4jNativePlatform  // Libs.nd4jCuda80Platform

To run tests from core (e.g. XorTest) place nd4j to your classpath and import nd4j implementation instead of jvm:

import nexus.impl.nd4j._
ctongfei commented 6 years ago

What do you think we should do for things like ArgMax if we use nd4j as a backend? The normal typing should be 

x: FloatTensor[A, B, C]
x |> ArgMaxAlong(B) // type: IntTensor[A, C]
ctongfei commented 5 years ago

Closed due to inactivity.