Merge pull request #141 from dfdx/fix/auto-seed-on-cuda

Fix device of seed in the :auto mode

Project.toml

@@ -1,7 +1,7 @@
 name = "Yota"
 uuid = "cd998857-8626-517d-b929-70ad188a48f0"
 authors = ["Andrei Zhabinski <andrei.zhabinski@gmail.com>"]
-version = "0.8.4"
+version = "0.8.5"
 
 [deps]
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"

src/grad.jl

@@ -132,8 +132,6 @@ with the following chain or calls:
 where `val = fn(args...)` and `pb` is the pullback function.
 """
 function chainrules_transform!(tape::Tape)
- # global TAPE = tape
- # error("")
  i = 1
  while i <= length(tape)
  # tape[V(i)] isa Call && tape[V(i)].fn == Core.kwcall && break
@@ -183,7 +181,6 @@ function step_back!(tape::Tape, y::Variable)
  end
  for (i, x) in enumerate(y_fargs)
  if x isa V
- global STATE = (tape, y, y_fargs, i, x)
  dx = push!(tape, mkcall(getfield, dxs, i; line="d$y/d$x"))
  # @debug "Updating derivative: $x -> $dx"
  set_or_add_deriv!(tape, x, dx)
@@ -208,8 +205,8 @@ function back!(tape::Tape; seed=1)
  error("Gradient of a vector-valued function requires a seed")
  elseif seed == :auto
  zval = tape[z].val
- # @assert zval isa Number || zval isa AbstractArray
- seed = zval isa AbstractArray ? ones(eltype(zval), size(zval)) : one(zval)
+ @assert zval isa Number || zval isa AbstractArray
+ seed = zval isa AbstractArray ? array_like(1, zval, size(zval)) : one(zval)
  end
  dy = push!(tape, Constant(seed; line="seed for $(tape[V(1)].val)"))
  # save seed var to use in compilation later

src/helpers.jl

@@ -88,7 +88,7 @@ end
 unbroadcast_prod_y(x::ArrayOrBroadcasted, y::ArrayOrBroadcasted, Δ) = unbroadcast_prod_x(y, x, Δ)
 
 # device_like(example, a) = (device = guess_device([example]); device(a))
-array_like(value, example) = fill!(similar(example, (1,)), value)
+array_like(value, example, sz=(1,)) = fill!(similar(example, sz), value)
 
 # unbroadcast_prod_x(x::Number, y::AbstractArray, Δ) = unbroadcast_prod_x(device_like(y, [x]), y, Δ)[1]
 unbroadcast_prod_x(x::Number, y::ArrayOrBroadcasted, Δ) = unbroadcast_prod_x(array_like(x, y), y, Δ)[1]

test/test_grad.jl

@@ -254,6 +254,17 @@ end
  val, g = grad(x -> 2x, [1.0, 2.0, 3.0]; seed=ones(3))
  @test val == [2.0, 4.0, 6.0]
  @test g == (ZeroTangent(), [2.0, 2.0, 2.0])
+
+ val, g = grad(x -> 2x, [1.0, 2.0, 3.0]; seed=:auto)
+ @test val == [2.0, 4.0, 6.0]
+ @test g == (ZeroTangent(), [2.0, 2.0, 2.0])
+
+ if CUDA.functional()
+ CUDA.allowscalar(false)
+ val, g = grad(x -> 2x, cu([1.0, 2.0, 3.0]); seed=:auto)
+ @test val == cu([2.0, 4.0, 6.0])
+ @test g == (ZeroTangent(), cu([2.0, 2.0, 2.0]))
+ end
 end