LuxDL / Lux.jl

Elegant & Performant Scientific Machine Learning in Julia
https://lux.csail.mit.edu/
MIT License
446 stars 50 forks source link

Rethinking `eltype` conversions in Adaptors #684

Open avik-pal opened 1 month ago

avik-pal commented 1 month ago

Currently, we forward calls to cu, mtl, etc., from the device adaptors. This is well-motivated because users often don't want to do FP64 ops on GPUs.

But there is a clear design problem. Let's say we want to convert to FP64 GPUArray, the current way is Array FP64 --> CuArray FP32 --> CuArray FP64.

We can solve this by associating an eltype with the adaptor:

  1. Default for CPUAdaptor: nothing, i.e., preserve eltype while converting
  2. For all other GPU Backends: FP32, i.e. convert to FP32.