Skip to content

Commit 0af5bec

Browse files
milanklmilankl
authored
Merge pull request #15 from pabvald/dims-standard-keyword-argument
Dims standard keyword argument
2 parents de6436e + e393be1 commit 0af5bec

File tree

4 files changed

+171
-107
lines changed

4 files changed

+171
-107
lines changed

‎README.md‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,20 @@ is round-to-nearest in linear space. Alternatively, a second argument can be eit
131131
See a derivation of this below.
132132
Decompression as with linear quantization via the `Array()` function.
133133

134+
#### Quantizing Along a Specific Dimension
135+
136+
In some cases, you may want to quantize an array along a specific dimension independently. This is useful when each slice along that dimension represents a separate dataset or when the range of values varies significantly across slices. You can achieve this by specifying the `dims` keyword argument in the `LinQuantArray` function.
137+
138+
For example, to quantize a 3D array along the second dimension:
139+
140+
```julia
141+
julia> L1 = LinQuantArray{UInt8}(A, dims=2)
142+
143+
julia> L2 = LogQuantArray{UInt8}(A, dims=2)
144+
```
145+
146+
Each element in the resulting vector `L` is a `LinQuantArray` that represents a slice of the original array along the specified dimension. This allows for independent quantization of each slice.
147+
134148
## Theory
135149

136150
### Linear quantization

‎src/linquantarrays.jl‎

Lines changed: 54 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ Quantise an array linearly into a LinQuantArray.
3131

3232
function LinQuantization(
3333
::Type{T},
34-
A::AbstractArray;
34+
A::AbstractArray{<:Real};
3535
extrema::Option{Tuple}=nothing
3636
) where {T<:Integer}
3737
# range of values in A
@@ -64,16 +64,62 @@ function LinQuantization(
6464
return LinQuantArray{T, ndims(Q)}(Q, Amin, Amax)
6565
end
6666

67-
function LinQuantArray{U}(A::AbstractArray{T,N}; extrema::Option{Tuple}=nothing) where {U<:Integer, T, N}
68-
LinQuantization(U, A; extrema=extrema)
67+
68+
"""
69+
LinQuantArray(TInteger, A; dims, extrema=nothing)
70+
71+
Linear quantization independently for every element along dimension`dims` in array `A`.
72+
73+
# Arguments
74+
- `TInteger`: the type of the quantised array
75+
- `A`: the array to quantise
76+
- `dims`: the dimension along which to quantise
77+
- `extrema`: the minimum and maximum of the range, defaults to `nothing`.
78+
79+
# Returns
80+
- a Vector{LinQuantArray} with the quantised array and the minimum and maximum of the original range.
81+
"""
82+
function LinQuantArray(
83+
::Type{TInteger},
84+
A::AbstractArray{T,N};
85+
dims::Int,
86+
extrema::Option{Tuple} = nothing
87+
) where {TInteger<:Integer,T,N}
88+
@assert dims <= N "Can't quantize a $N-dimensional array in dim=$dims"
89+
n = size(A)[dims]
90+
L = Vector{LinQuantArray}(undef, n)
91+
t = [if j == dims 1 else Colon() end for j in 1:N]
92+
for i in 1:n
93+
t[dims] = i
94+
L[i] = LinQuantization(TInteger,A[t...]; extrema=extrema)
95+
end
96+
return L
97+
end
98+
99+
100+
function LinQuantArray{U}(
101+
A::AbstractArray{T,N};
102+
dims::Option{Int}=nothing,
103+
extrema::Option{Tuple}=nothing
104+
) where {U<:Integer,T,N}
105+
isnothing(dims) ? LinQuantization(U,A;extrema=extrema) : LinQuantArray(U,A;dims=dims,extrema=extrema)
106+
end
107+
108+
function LinQuant8Array(A::AbstractArray{T,N}; dims::Option{Int}=nothing) where {T,N}
109+
isnothing(dims) ? LinQuantization(UInt8,A) : LinQuantArray(UInt8,A; dims=dims)
110+
end
111+
112+
function LinQuant16Array(A::AbstractArray{T,N}; dims::Option{Int}=nothing) where {T,N}
113+
isnothing(dims) ? LinQuantization(UInt16,A) : LinQuantArray(UInt16,A; dims=dims)
69114
end
70115

71-
# keep compatibility: shortcuts for unsigned integers of 8, 16, 24 and 32 bit
72-
LinQuant8Array(A::AbstractArray{T,N}) where {T,N} = LinQuantization(UInt8,A)
73-
LinQuant16Array(A::AbstractArray{T,N}) where {T,N} = LinQuantization(UInt16,A)
74-
LinQuant24Array(A::AbstractArray{T,N}) where {T,N} = LinQuantization(UInt24,A)
75-
LinQuant32Array(A::AbstractArray{T,N}) where {T,N} = LinQuantization(UInt32,A)
116+
function LinQuant24Array(A::AbstractArray{T,N}; dims::Option{Int}=nothing) where {T,N}
117+
isnothing(dims) ? LinQuantization(UInt24,A) : LinQuantArray(UInt24,A; dims=dims)
118+
end
76119

120+
function LinQuant32Array(A::AbstractArray{T,N}; dims::Option{Int}=nothing) where {T,N}
121+
isnothing(dims) ? LinQuantization(UInt32,A) : LinQuantArray(UInt32,A; dims=dims)
122+
end
77123

78124
"""
79125
Array{U}(Q::LinQuantArray) where {U<:AbstractFloat}
@@ -119,51 +165,6 @@ Base.Array(Q::LinQuantArray{Int24,N}) where N = Array{Float32}(Q)
119165
Base.Array(Q::LinQuantArray{Int32,N}) where N = Array{Float64}(Q)
120166

121167

122-
"""
123-
LinQuantArray(TInteger, A, dim; extrema=nothing)
124-
125-
Linear quantization independently for every element along dimension
126-
dim in array A.
127-
128-
# Arguments
129-
- `TInteger`: the type of the quantised array
130-
- `A`: the array to quantise
131-
- `dim`: the dimension along which to quantise
132-
- `extrema`: the minimum and maximum of the range, defaults to `nothing`.
133-
134-
# Returns
135-
- a Vector{LinQuantArray} with the quantised array and the minimum and maximum of the original range.
136-
"""
137-
function LinQuantArray(
138-
::Type{TInteger},
139-
A::AbstractArray{T,N},
140-
dim::Integer;
141-
extrema::Option{Tuple} = nothing
142-
) where {TInteger,T,N}
143-
@assert dim <= N "Can't quantize a $N-dimensional array in dim=$dim"
144-
n = size(A)[dim]
145-
L = Vector{LinQuantArray}(undef, n)
146-
t = [if j == dim 1 else Colon() end for j in 1:N]
147-
for i in 1:n
148-
t[dim] = i
149-
L[i] = LinQuantization(TInteger,A[t...]; extrema=extrema)
150-
end
151-
return L
152-
end
153-
154-
function LinQuantArray{U}(
155-
A::AbstractArray{T,N},
156-
dim::Int;
157-
extrema::Option{Tuple}=nothing
158-
) where {U<:Integer,T,N}
159-
LinQuantArray(U,A,dim;extrema=extrema)
160-
end
161-
162-
# keep compatibility: shortcuts for unsigned integers of 8, 16, 24 and 32 bit
163-
LinQuant8Array(A::AbstractArray{T,N},dim::Int) where {T,N} = LinQuantArray(UInt8,A,dim)
164-
LinQuant16Array(A::AbstractArray{T,N},dim::Int) where {T,N} = LinQuantArray(UInt16,A,dim)
165-
LinQuant24Array(A::AbstractArray{T,N},dim::Int) where {T,N} = LinQuantArray(UInt24,A,dim)
166-
LinQuant32Array(A::AbstractArray{T,N},dim::Int) where {T,N} = LinQuantArray(UInt32,A,dim)
167168

168169
"""
169170
Array{U}(L::Vector{LinQuantArray}) where {U<:AbstractFloat}

‎src/logquantarrays.jl‎

Lines changed: 71 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ struct LogQuantArray{T,N} <: AbstractArray{Unsigned,N}
1010
end
1111

1212
Base.size(QA::LogQuantArray) = size(QA.A)
13-
Base.getindex(QA::LogQuantArray,i...) = getindex(QA.A,i...)
13+
Base.getindex(QA::LogQuantArray, i...) = getindex(QA.A, i...)
1414
Base.eltype(Q::LogQuantArray{T,N}) where {T,N} = T
1515

1616
"""
@@ -21,7 +21,7 @@ Base.eltype(Q::LogQuantArray{T,N}) where {T,N} = T
2121
"""
2222
function minpos(A::AbstractArray{T}) where T
2323
o = zero(T)
24-
mi = foldl((x,y) -> y > o ? min(x,y) : x, A; init=typemax(T))
24+
mi = foldl((x, y) -> y > o ? min(x, y) : x, A; init=typemax(T))
2525
mi == typemax(T) && return o
2626
return mi
2727
end
@@ -39,7 +39,7 @@ Quantize elements of an array logarithmically into UInts with either round to ne
3939
# Returns
4040
- a LogQuantArray{T} with the quantised array and the minimum and maximum of the original range.
4141
"""
42-
function LogQuantization(
42+
function LogQuantization(
4343
::Type{T},
4444
A::AbstractArray,
4545
round_nearest_in::Symbol=:linspace
@@ -60,38 +60,86 @@ function LogQuantization(
6060
else
6161
# inverse log spacing
6262
# map min to 1 and max to ff..., reserve 0 for 0.
63-
Δ⁻¹ = (2^(sizeof(T)*8)-2)/(logmax-logmin)
64-
63+
Δ⁻¹ = (2^(sizeof(T) * 8) - 2) / (logmax - logmin)
64+
6565
# shift to round-to-nearest in lin or log-space
6666
if round_nearest_in == :linspace
67-
c = 1/2 - Δ⁻¹*log(mi*(exp(1/Δ⁻¹)+1)/2)
67+
c = 1 / 2 - Δ⁻¹ * log(mi * (exp(1 / Δ⁻¹) + 1) / 2)
6868
elseif round_nearest_in == :logspace
69-
c = -logmin*Δ⁻¹
69+
c = -logmin * Δ⁻¹
7070
else
7171
throw(ArgumentError("Round-to-nearest either :linspace or :logspace"))
7272
end
7373
end
7474

7575
# preallocate output
76-
Q = similar(A,T)
76+
Q = similar(A, T)
7777

7878
@inbounds for i in eachindex(A)
7979
# store 0 as 0x00...
8080
# store positive numbers via convert to logpacking as 0x1-0xff..
81-
Q[i] = iszero(A[i]) ? zero(T) : convert(T,round(c + Δ⁻¹*log(Float64(A[i]))))+one(T)
81+
Q[i] = iszero(A[i]) ? zero(T) : convert(T, round(c + Δ⁻¹ * log(Float64(A[i])))) + one(T)
82+
end
83+
84+
return LogQuantArray{T,ndims(Q)}(Q, Float64(logmin), Float64(logmax))
85+
end
86+
87+
88+
89+
"""
90+
LogQuantArray(::Type{TUInt},A::AbstractArray{T,N};dims::Int) where {TUInt<:Unsigned,T,N}
91+
Logarithmic quantization independently for every element along dimension
92+
`dims` in array `A`.
93+
94+
# Returns
95+
- a Vector{LogQuantArray}.
96+
"""
97+
function LogQuantArray(::Type{TUInt}, A::AbstractArray{T,N}; dims::Int) where {TUInt<:Unsigned,T,N}
98+
@assert dims <= N "Can't quantize a $N-dimensional array in dimension=$dims"
99+
n = size(A)[dims]
100+
L = Vector{LogQuantArray}(undef, n)
101+
t = [
102+
if j == dims
103+
1
104+
else
105+
Colon()
106+
end for j in 1:N
107+
]
108+
for i in 1:n
109+
t[dims] = i
110+
L[i] = LogQuantization(TUInt, A[t...])
82111
end
112+
return L
113+
end
83114

84-
return LogQuantArray{T,ndims(Q)}(Q,Float64(logmin),Float64(logmax))
115+
function LogQuantArray{U}(
116+
A::AbstractArray{T,N};
117+
dims::Option{Int}=nothing
118+
) where {U<:Unsigned,T,N}
119+
isnothing(dims) ? LogQuantization(U,A) : LogQuantArray(U,A;dims=dims)
85120
end
86121

87-
# define for 8, 16, 24 and 32 bit uints
88-
LogQuant8Array(A::AbstractArray{T,N},rn::Symbol=:linspace) where {T,N} = LogQuantization(UInt8,A,rn)
89-
LogQuant16Array(A::AbstractArray{T,N},rn::Symbol=:linspace) where {T,N} = LogQuantization(UInt16,A,rn)
90-
LogQuant24Array(A::AbstractArray{T,N},rn::Symbol=:linspace) where {T,N} = LogQuantization(UInt24,A,rn)
91-
LogQuant32Array(A::AbstractArray{T,N},rn::Symbol=:linspace) where {T,N} = LogQuantization(UInt32,A,rn)
122+
# keep compatibility: shortcuts for unsigned integers of 8, 16, 24 and 32-bit
123+
function LogQuant8Array(A::AbstractArray{T,N}, rn::Symbol=:linspace; dims::Option{Int}=nothing) where {T,N}
124+
isnothing(dims) ? LogQuantization(UInt8, A, rn) : LogQuantArray(UInt8, A; dims=dims)
125+
end
126+
127+
function LogQuant16Array(A::AbstractArray{T,N}, rn::Symbol=:linspace; dims::Option{Int}=nothing) where {T,N}
128+
isnothing(dims) ? LogQuantization(UInt16, A, rn) : LogQuantArray(UInt16, A; dims=dims)
129+
end
130+
131+
function LogQuant24Array(A::AbstractArray{T,N}, rn::Symbol=:linspace; dims::Option{Int}=nothing) where {T,N}
132+
isnothing(dims) ? LogQuantization(UInt24, A, rn) : LogQuantArray(UInt24, A; dims=dims)
133+
end
134+
135+
function LogQuant32Array(A::AbstractArray{T,N}, rn::Symbol=:linspace; dims::Option{Int}=nothing) where {T,N}
136+
isnothing(dims) ? LogQuantization(UInt32, A, rn) : LogQuantArray(UInt32, A; dims=dims)
137+
end
138+
139+
92140

93141
"""
94-
Array{T}(n::Integer,Q::LogQuantArray) where {T<:AbstractFloat}
142+
Array{T}(n::Integer, Q::LogQuantArray) where {T<:AbstractFloat}
95143
96144
De-quantise a LogQuantArray into floats.
97145
@@ -102,19 +150,19 @@ De-quantise a LogQuantArray into floats.
102150
# Returns
103151
- an array of type T with the de-quantised values.
104152
"""
105-
function Base.Array{T}(n::Integer,Q::LogQuantArray) where {T<:AbstractFloat}
153+
function Base.Array{T}(n::Integer, Q::LogQuantArray) where {T<:AbstractFloat}
106154
Qlogmin = Q.min # log(min::Float64)
107155
Qlogmax = Q.max # log(max::Float64)
108156

109157
# spacing in logspace ::Float64
110-
Δ = (Qlogmax-Qlogmin)/(2^n-2) # -2 as 0x00.. is reserved for 0
158+
Δ = (Qlogmax - Qlogmin) / (2^n - 2) # -2 as 0x00.. is reserved for 0
111159

112-
A = similar(Q,T) # preallocate
160+
A = similar(Q, T) # preallocate
113161

114162
@inbounds for i in eachindex(A)
115163
# 0x0 is unpack as 0
116164
# exp in Float64 then convert to T at assignment =
117-
A[i] = iszero(Q[i]) ? zero(T) : A[i] = exp(Qlogmin + (Q[i]-1)*Δ)
165+
A[i] = iszero(Q[i]) ? zero(T) : A[i] = exp(Qlogmin + (Q[i] - 1) * Δ)
118166
end
119167

120168
return A
@@ -179,9 +227,9 @@ function Base.Array{T}(L::Vector{LogQuantArray}) where T
179227
n = length(L)
180228
s = size(L[1])
181229
t = axes(L[1])
182-
A = Array{T,N+1}(undef,s...,length(L))
230+
A = Array{T,N + 1}(undef, s..., length(L))
183231
for i in 1:n
184-
A[t...,i] = Array{T}(L[i])
232+
A[t..., i] = Array{T}(L[i])
185233
end
186234
return A
187-
end
235+
end

0 commit comments

Comments
 (0)