Skip to content

Commit

Permalink
more general indexing
Browse files Browse the repository at this point in the history
  • Loading branch information
CarloLucibello committed Dec 15, 2017
1 parent 7d14a44 commit fe1e45a
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 3 deletions.
19 changes: 16 additions & 3 deletions src/karray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -494,9 +494,6 @@ end
# function _getindex(l::LinearIndexing, A::AbstractArray, I::Union{Real, AbstractArray, Colon}...)
# in abstractarray.jl:487,multidimensional.jl:184.

if VERSION < v"0.5.0"
@typealias6 AbstractUnitRange UnitRange
end

function getindex{T}(A::KnetArray{T}, I::AbstractUnitRange)
if !(1 <= first(I) <= last(I) <= length(A)); throw(BoundsError(A,I)); end
Expand Down Expand Up @@ -552,6 +549,22 @@ function setindex!{T}(A::KnetArray{T}, v, I::Colon)
unsafe_copy!(A,1,v,1,length(A))
end

## General Indexing Fallback to linear indexing

function getindex(a::KnetArray, I...)
crange = CartesianRange(to_indices(a, I))
linind = [sub2ind(size(a), t.I...) for t in crange]
b = getindex(a, vec(linind))
shape = size(crange) # TODO drop scalar dimension
reshape(b, shape)
end

function setindex!(a::KnetArray, v, I...)
crange = CartesianRange(to_indices(a, I))
linind = [sub2ind(size(a), t.I...) for t in crange]
setindex!(a, v, vec(linind))
end

for F in (32,64); T=Symbol("Float$F"); @eval begin

## Indexing with KnetArray{Int32}: low level, only Int32 supported, no bounds checking
Expand Down
20 changes: 20 additions & 0 deletions test/karray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,26 @@ if gpu() >= 0
@test size(reshape(a, :, 4)) == size(reshape(a, (:, 4))) == (2, 4)
@test size(reshape(a, :, 1, 4)) == (2, 1, 4)
end

@testset "general indexing" begin
# TODO
#
#julia> a=KnetArray(rand(3,3,2));
#
#julia> grad(a->sum(a[1:2,:,1]))(a)
#3×3×2 Knet.KnetArray{Float64,3}:
#[:, :, 1] =
# 1.0 1.0 1.0
# 1.0 1.0 1.0
# 0.0 0.0 0.0
#
#[:, :, 2] =
# 0.0 0.0 0.0
# 0.0 0.0 0.0
# 0.0 0.0 0.0
end


end
end

Expand Down

0 comments on commit fe1e45a

Please sign in to comment.