Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Code suggestions for #1210 #1213

Merged
merged 44 commits into from
Jan 2, 2024
Merged

Code suggestions for #1210 #1213

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
44 commits
Select commit Hold shift + click to select a range
8b235b9
move `ex.args[2] isa Integer`
hyrodium Oct 29, 2023
477ab09
split `if` block
hyrodium Oct 29, 2023
684bd0a
simplify :zeros and :ones
hyrodium Oct 29, 2023
4885ad9
refactor :rand
hyrodium Oct 29, 2023
d048d16
refactor :randn and :randexp
hyrodium Oct 29, 2023
835ad76
update comments
hyrodium Oct 29, 2023
5cbc347
add _isnonnegvec
hyrodium Oct 30, 2023
d279187
update with `_isnonnegvec`
hyrodium Oct 30, 2023
2ed9633
add `_isnonnegvec(args, n)` method to check the size of `args`
hyrodium Oct 31, 2023
16dbbc7
fix `@SArray` for `@SArray rand(rng,T,dim)` etc.
hyrodium Oct 31, 2023
563ea0c
update comments
hyrodium Oct 31, 2023
d1a2d32
update `@SVector` macro
hyrodium Oct 31, 2023
b3cf411
update `@SMatrix`
hyrodium Nov 1, 2023
ff9b32a
update `@SVector`
hyrodium Nov 1, 2023
6eb7829
update `@SArray`
hyrodium Nov 1, 2023
0a9fc2b
introduce `fargs` variable
hyrodium Nov 12, 2023
337a2e0
avoid `_isnonnegvec` in `static_matrix_gen`
hyrodium Nov 12, 2023
5021855
avoid `_isnonnegvec` in `static_vector_gen`
hyrodium Nov 12, 2023
7352e26
remove unnecessary `_isnonnegvec`
hyrodium Nov 18, 2023
623cb0c
add `_rng()` function
hyrodium Nov 18, 2023
d1cf08d
update tests on `@SVector` macro
hyrodium Nov 18, 2023
0403edf
update tests on `@MVector` macro
hyrodium Dec 1, 2023
9f0753b
organize test/MMatrix.jl and test/SMatrix.jl
hyrodium Dec 2, 2023
12b7634
organize test/MMatrix.jl and test/SMatrix.jl
hyrodium Dec 2, 2023
cb07fcf
update with broken tests
hyrodium Dec 2, 2023
b4ace0d
organize test/MMatrix.jl and test/SMatrix.jl for `rand*` functions
hyrodium Dec 2, 2023
6c5a02b
fix around `broken` key for `@test` macro
hyrodium Dec 2, 2023
9217564
fix zero-length tests
hyrodium Dec 2, 2023
1b0c4f4
update `test/SArray.jl` to match `test/MArray.jl`
hyrodium Jan 2, 2024
24265b9
update tests for `@SArray ones` etc.
hyrodium Jan 2, 2024
c9df6a8
add supports for `@SArray ones(3-1,2)` etc.
hyrodium Jan 2, 2024
929e692
move block for `fill`
hyrodium Jan 2, 2024
4feb7af
update macro `@SArray rand(rng,2,3)` to use ordinary dispatches
hyrodium Jan 2, 2024
d17ee6b
update around `@SArray randn` etc.
hyrodium Jan 2, 2024
af095e3
remove unnecessary dollars
hyrodium Jan 2, 2024
26c922e
simplify `@SArray fill`
hyrodium Jan 2, 2024
2cbb67a
add `@testset "expand_error"`
hyrodium Jan 2, 2024
6a03236
update tests for `@SArray rand(...)` etc.
hyrodium Jan 2, 2024
0f1c559
fix bug in `rand*_with_Val`
hyrodium Jan 2, 2024
89fc5e5
cleanup tests
hyrodium Jan 2, 2024
c55b170
update macro `@SMatrix rand(rng,2,3)` to use ordinary dispatches
hyrodium Jan 2, 2024
0bb47e7
update macro `@SVector rand(rng,3)` to use ordinary dispatches
hyrodium Jan 2, 2024
c8cb1f5
move block for `fill`
hyrodium Jan 2, 2024
2e40f73
simplify `_randexp_with_Val`
hyrodium Jan 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
142 changes: 93 additions & 49 deletions src/SArray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,22 +142,65 @@ function parse_cat_ast(ex::Expr)
cat_any(Val(maxdim), Val(catdim), nargs)
end

#=
For example,
* `@SArray rand(2, 3, 4)`
* `@SArray rand(rng, 3, 4)`
will be expanded to the following.
* `_rand_with_Val(SArray, 2, 3, _int2val(2), _int2val(3), Val((4,)))`
* `_rand_with_Val(SArray, 2, 3, _int2val(rng), _int2val(3), Val((4,)))`
The function `_int2val` is required to avoid the following case.
* `_rand_with_Val(SArray, 2, 3, Val(2), Val(3), Val((4,)))`
* `_rand_with_Val(SArray, 2, 3, Val(rng), Val(3), Val((4,)))`
Mutable object such as `rng` cannot be type parameter, and `Val(rng)` throws an error.
=#
_int2val(x::Int) = Val(x)
_int2val(::Any) = nothing
# @SArray zeros(...)
_zeros_with_Val(::Type{SA}, ::Int, ::Val{n1}, ::Val{ns}) where {SA, n1, ns} = zeros(SA{Tuple{n1, ns...}})
_zeros_with_Val(::Type{SA}, T::DataType, ::Val, ::Val{ns}) where {SA, ns} = zeros(SA{Tuple{ns...}, T})
# @SArray ones(...)
_ones_with_Val(::Type{SA}, ::Int, ::Val{n1}, ::Val{ns}) where {SA, n1, ns} = ones(SA{Tuple{n1, ns...}})
_ones_with_Val(::Type{SA}, T::DataType, ::Val, ::Val{ns}) where {SA, ns} = ones(SA{Tuple{ns...}, T})
# @SArray rand(...)
_rand_with_Val(::Type{SA}, ::Int, ::Int, ::Val{n1}, ::Val{n2}, ::Val{ns}) where {SA, n1, n2, ns} = rand(SA{Tuple{n1,n2,ns...}})
_rand_with_Val(::Type{SA}, T::DataType, ::Int, ::Nothing, ::Val{n1}, ::Val{ns}) where {SA, n1, ns} = _rand(Random.GLOBAL_RNG, T, Size(n1, ns...), SA{Tuple{n1, ns...}, T})
_rand_with_Val(::Type{SA}, sampler, ::Int, ::Nothing, ::Val{n1}, ::Val{ns}) where {SA, n1, ns} = _rand(Random.GLOBAL_RNG, sampler, Size(n1, ns...), SA{Tuple{n1, ns...}, Random.gentype(sampler)})
_rand_with_Val(::Type{SA}, rng::AbstractRNG, ::Int, ::Nothing, ::Val{n1}, ::Val{ns}) where {SA, n1, ns} = _rand(rng, Float64, Size(n1, ns...), SA{Tuple{n1, ns...}, Float64})
_rand_with_Val(::Type{SA}, rng::AbstractRNG, T::DataType, ::Nothing, ::Nothing, ::Val{ns}) where {SA, ns} = _rand(rng, T, Size(ns...), SA{Tuple{ns...}, T})
_rand_with_Val(::Type{SA}, rng::AbstractRNG, sampler, ::Nothing, ::Nothing, ::Val{ns}) where {SA, ns} = _rand(rng, sampler, Size(ns...), SA{Tuple{ns...}, Random.gentype(sampler)})
# @SArray randn(...)
_randn_with_Val(::Type{SA}, ::Int, ::Int, ::Val{n1}, ::Val{n2}, ::Val{ns}) where {SA, n1, n2, ns} = randn(SA{Tuple{n1,n2,ns...}})
_randn_with_Val(::Type{SA}, T::DataType, ::Int, ::Nothing, ::Val{n1}, ::Val{ns}) where {SA, n1, ns} = _randn(Random.GLOBAL_RNG, Size(n1, ns...), SA{Tuple{n1, ns...}, T})
_randn_with_Val(::Type{SA}, rng::AbstractRNG, ::Int, ::Nothing, ::Val{n1}, ::Val{ns}) where {SA, n1, ns} = _randn(rng, Size(n1, ns...), SA{Tuple{n1, ns...}, Float64})
_randn_with_Val(::Type{SA}, rng::AbstractRNG, T::DataType, ::Nothing, ::Nothing, ::Val{ns}) where {SA, ns} = _randn(rng, Size(ns...), SA{Tuple{ns...}, T})
# @SArray randexp(...)
_randexp_with_Val(::Type{SA}, ::Int, ::Int, ::Val{n1}, ::Val{n2}, ::Val{ns}) where {SA, n1, n2, ns} = randexp(SA{Tuple{n1,n2,ns...}})
_randexp_with_Val(::Type{SA}, T::DataType, ::Int, ::Nothing, ::Val{n1}, ::Val{ns}) where {SA, n1, ns} = _randexp(Random.GLOBAL_RNG, Size(n1, ns...), SA{Tuple{n1, ns...}, T})
_randexp_with_Val(::Type{SA}, rng::AbstractRNG, ::Int, ::Nothing, ::Val{n1}, ::Val{ns}) where {SA, n1, ns} = _randexp(rng, Size(n1, ns...), SA{Tuple{n1, ns...}, Float64})
_randexp_with_Val(::Type{SA}, rng::AbstractRNG, T::DataType, ::Nothing, ::Nothing, ::Val{ns}) where {SA, ns} = _randexp(rng, Size(ns...), SA{Tuple{ns...}, T})

escall(args) = Iterators.map(esc, args)
function _isnonnegvec(args)
length(args) == 0 && return false
all(isa.(args, Integer)) && return all(args .≥ 0)
return false
end
function static_array_gen(::Type{SA}, @nospecialize(ex), mod::Module) where {SA}
if !isa(ex, Expr)
error("Bad input for @$SA")
end
head = ex.head
if head === :vect # vector
return :($SA{$Tuple{$(length(ex.args))}}($tuple($(escall(ex.args)...))))
return :($SA{Tuple{$(length(ex.args))}}($tuple($(escall(ex.args)...))))
elseif head === :ref # typed, vector
return :($SA{$Tuple{$(length(ex.args)-1)},$(esc(ex.args[1]))}($tuple($(escall(ex.args[2:end])...))))
return :($SA{Tuple{$(length(ex.args)-1)},$(esc(ex.args[1]))}($tuple($(escall(ex.args[2:end])...))))
elseif head === :typed_vcat || head === :typed_hcat || head === :typed_ncat # typed, cat
args = parse_cat_ast(ex)
return :($SA{$Tuple{$(size(args)...)},$(esc(ex.args[1]))}($tuple($(escall(args)...))))
return :($SA{Tuple{$(size(args)...)},$(esc(ex.args[1]))}($tuple($(escall(args)...))))
elseif head === :vcat || head === :hcat || head === :ncat # untyped, cat
args = parse_cat_ast(ex)
return :($SA{$Tuple{$(size(args)...)}}($tuple($(escall(args)...))))
return :($SA{Tuple{$(size(args)...)}}($tuple($(escall(args)...))))
elseif head === :comprehension
if length(ex.args) != 1
error("Expected generator in comprehension, e.g. [f(i,j) for i = 1:3, j = 1:3]")
Expand All @@ -173,7 +216,7 @@ function static_array_gen(::Type{SA}, @nospecialize(ex), mod::Module) where {SA}
return quote
let
f($(escall(rng_args)...)) = $(esc(ex.args[1]))
$SA{$Tuple{$(size(exprs)...)}}($tuple($(exprs...)))
$SA{Tuple{$(size(exprs)...)}}($tuple($(exprs...)))
end
end
elseif head === :typed_comprehension
Expand All @@ -192,57 +235,58 @@ function static_array_gen(::Type{SA}, @nospecialize(ex), mod::Module) where {SA}
return quote
let
f($(escall(rng_args)...)) = $(esc(ex.args[1]))
$SA{$Tuple{$(size(exprs)...)},$T}($tuple($(exprs...)))
$SA{Tuple{$(size(exprs)...)},$T}($tuple($(exprs...)))
end
end
elseif head === :call
f = ex.args[1]
if f === :zeros || f === :ones || f === :rand || f === :randn || f === :randexp
if length(ex.args) == 1
f === :zeros || f === :ones || error("@$SA got bad expression: $(ex)")
return :($f($SA{$Tuple{},$Float64}))
elseif f !== :rand || length(ex.args) == 2
return quote
if isa($(esc(ex.args[2])), DataType)
$f($SA{$Tuple{$(escall(ex.args[3:end])...)},$(esc(ex.args[2]))})
else
$f($SA{$Tuple{$(escall(ex.args[2:end])...)}})
end
end
fargs = ex.args[2:end]
if f === :zeros || f === :ones
_f_with_Val = Symbol(:_, f, :_with_Val)
if length(fargs) == 0
# for calls like `zeros()`
return :($f($SA{Tuple{},$Float64}))
elseif _isnonnegvec(fargs)
# for calls like `zeros(dims...)`
return :($f($SA{Tuple{$(escall(fargs)...)}}))
else
return quote
if isa($(esc(ex.args[2])), DataType)
$f($SA{$Tuple{$(escall(ex.args[3:end])...)},$(esc(ex.args[2]))})
elseif isa($(esc(ex.args[2])), Integer)
$f($SA{$Tuple{$(escall(ex.args[2:end])...)}})
elseif isa($(esc(ex.args[2])), Random.AbstractRNG)
# for calls like rand(rng::AbstractRNG, sampler, dims::Integer...)
StaticArrays._rand(
$(esc(ex.args[2])),
$(esc(ex.args[3])),
Size($(escall(ex.args[4:end])...)),
$SA{
Tuple{$(escall(ex.args[4:end])...)},
Random.gentype($(esc(ex.args[3]))),
},
)
else
# for calls like rand(sampler, dims::Integer...)
StaticArrays._rand(
Random.GLOBAL_RNG,
$(esc(ex.args[2])),
Size($(escall(ex.args[3:end])...)),
$SA{
Tuple{$(escall(ex.args[3:end])...)},
Random.gentype($(esc(ex.args[2]))),
},
)
end
end
# for calls like `zeros(type)`
# for calls like `zeros(type, dims...)`
return :($_f_with_Val($SA, $(esc(fargs[1])), Val($(esc(fargs[1]))), Val(tuple($(escall(fargs[2:end])...)))))
end
elseif f === :fill
length(ex.args) == 1 && error("@$SA got bad expression: $(ex)")
return :($f($(esc(ex.args[2])), $SA{$Tuple{$(escall(ex.args[3:end])...)}}))
# for calls like `fill(value, dims...)`
return :($f($(esc(fargs[1])), $SA{Tuple{$(escall(fargs[2:end])...)}}))
elseif f === :rand || f === :randn || f === :randexp
_f_with_Val = Symbol(:_, f, :_with_Val)
if length(fargs) == 0
# No support for `@SArray rand()`
error("@$SA got bad expression: $(ex)")
elseif _isnonnegvec(fargs)
# for calls like `rand(dims...)`
return :($f($SA{Tuple{$(escall(fargs)...)}}))
elseif length(fargs) ≥ 2
# for calls like `rand(dim1, dim2, dims...)`
# for calls like `rand(type, dim1, dims...)`
# for calls like `rand(sampler, dim1, dims...)`
# for calls like `rand(rng, dim1, dims...)`
# for calls like `rand(rng, type, dims...)`
# for calls like `rand(rng, sampler, dims...)`
# for calls like `randn(dim1, dim2, dims...)`
# for calls like `randn(type, dim1, dims...)`
# for calls like `randn(rng, dim1, dims...)`
# for calls like `randn(rng, type, dims...)`
# for calls like `randexp(dim1, dim2, dims...)`
# for calls like `randexp(type, dim1, dims...)`
# for calls like `randexp(rng, dim1, dims...)`
# for calls like `randexp(rng, type, dims...)`
return :($_f_with_Val($SA, $(esc(fargs[1])), $(esc(fargs[2])), _int2val($(esc(fargs[1]))), _int2val($(esc(fargs[2]))), Val(tuple($(escall(fargs[3:end])...)))))
elseif length(fargs) == 1
# for calls like `rand(dim)`
return :($f($SA{Tuple{$(escall(fargs)...)}}))
else
error("@$SA got bad expression: $(ex)")
end
else
error("@$SA only supports the zeros(), ones(), fill(), rand(), randn(), and randexp() functions.")
end
Expand Down
86 changes: 54 additions & 32 deletions src/SMatrix.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,21 @@ function check_matrix_size(x::Tuple, T = :S)
x1, x2
end

# @SMatrix rand(...)
_rand_with_Val(::Type{SM}, rng::AbstractRNG, ::Val{n1}, ::Val{n2}) where {SM, n1, n2} = rand(rng, SM{n1, n2})
_rand_with_Val(::Type{SM}, T::DataType, ::Val{n1}, ::Val{n2}) where {SM, n1, n2} = _rand(Random.GLOBAL_RNG, T, Size(n1, n2), SM{n1, n2, T})
_rand_with_Val(::Type{SM}, sampler, ::Val{n1}, ::Val{n2}) where {SM, n1, n2} = _rand(Random.GLOBAL_RNG, sampler, Size(n1, n2), SM{n1, n2, Random.gentype(sampler)})
_rand_with_Val(::Type{SM}, rng::AbstractRNG, T::DataType, ::Val{n1}, ::Val{n2}) where {SM, n1, n2} = rand(rng, SM{n1, n2, T})
_rand_with_Val(::Type{SM}, rng::AbstractRNG, sampler, ::Val{n1}, ::Val{n2}) where {SM, n1, n2} = _rand(rng, sampler, Size(n1, n2), SM{n1, n2, Random.gentype(sampler)})
# @SMatrix randn(...)
_randn_with_Val(::Type{SM}, rng::AbstractRNG, ::Val{n1}, ::Val{n2}) where {SM, n1, n2} = randn(rng, SM{n1, n2})
_randn_with_Val(::Type{SM}, T::DataType, ::Val{n1}, ::Val{n2}) where {SM, n1, n2} = _randn(Random.GLOBAL_RNG, Size(n1, n2), SM{n1, n2, T})
_randn_with_Val(::Type{SM}, rng::AbstractRNG, T::DataType, ::Val{n1}, ::Val{n2}) where {SM, n1, n2} = randn(rng, SM{n1, n2, T})
# @SMatrix randexp(...)
_randexp_with_Val(::Type{SM}, rng::AbstractRNG, ::Val{n1}, ::Val{n2}) where {SM, n1, n2} = randexp(rng, SM{n1, n2})
_randexp_with_Val(::Type{SM}, T::DataType, ::Val{n1}, ::Val{n2}) where {SM, n1, n2} = _randexp(Random.GLOBAL_RNG, Size(n1, n2), SM{n1, n2, T})
_randexp_with_Val(::Type{SM}, rng::AbstractRNG, T::DataType, ::Val{n1}, ::Val{n2}) where {SM, n1, n2} = randexp(rng, SM{n1, n2, T})

function static_matrix_gen(::Type{SM}, @nospecialize(ex), mod::Module) where {SM}
if !isa(ex, Expr)
error("Bad input for @$SM")
Expand Down Expand Up @@ -69,44 +84,51 @@ function static_matrix_gen(::Type{SM}, @nospecialize(ex), mod::Module) where {SM
end
elseif head === :call
f = ex.args[1]
if f === :zeros || f === :ones || f === :rand || f === :randn || f === :randexp
if length(ex.args) == 3
return :($f($SM{$(escall(ex.args[2:3])...), Float64})) # default to Float64 like Base
elseif length(ex.args) == 4
if f === :rand && ex.args[3] isa Int && ex.args[3] ≥ 0 && ex.args[4] isa Int && ex.args[4] ≥ 0
# for calls like rand(sampler, n, m) or rand(type, n, m)
return quote
StaticArrays._rand(
Random.GLOBAL_RNG,
$(esc(ex.args[2])),
Size($(esc(ex.args[3])), $(esc(ex.args[4]))),
$SM{$(esc(ex.args[3])), $(esc(ex.args[4])), Random.gentype($(esc(ex.args[2])))},
)
end
else
return :($f($SM{$(escall(ex.args[[3,4,2]])...)}))
end
elseif length(ex.args) == 5 && f === :rand && ex.args[4] isa Int && ex.args[4] ≥ 0 && ex.args[5] isa Int && ex.args[5] ≥ 0
# for calls like rand(rng::AbstractRNG, sampler, n, m) or rand(rng::AbstractRNG, type, n, m)
return quote
StaticArrays._rand(
$(esc(ex.args[2])),
$(esc(ex.args[3])),
Size($(esc(ex.args[4])), $(esc(ex.args[5]))),
$SM{$(esc(ex.args[4])), $(esc(ex.args[5])), Random.gentype($(esc(ex.args[3])))},
)
end
fargs = ex.args[2:end]
if f === :zeros || f === :ones
if length(fargs) == 2
# for calls like `zeros(dim1, dim2)`
return :($f($SM{$(escall(fargs)...)}))
elseif length(fargs[2:end]) == 2
# for calls like `zeros(type, dim1, dim2)`
return :($f($SM{$(escall(fargs[2:end])...), $(esc(fargs[1]))}))
else
error("@$SM expected a 2-dimensional array expression")
error("@$SM got bad expression: $(ex)")
end
elseif ex.args[1] === :fill
if length(ex.args) == 4
return :($f($(esc(ex.args[2])), $SM{$(escall(ex.args[3:4])...)}))
elseif f === :fill
# for calls like `fill(value, dim1, dim2)`
if length(fargs[2:end]) == 2
return :($f($(esc(fargs[1])), $SM{$(escall(fargs[2:end])...)}))
else
error("@$SM expected a 2-dimensional array expression")
end
elseif f === :rand || f === :randn || f === :randexp
_f_with_Val = Symbol(:_, f, :_with_Val)
if length(fargs) == 2
# for calls like `rand(dim1, dim2)`
# for calls like `randn(dim1, dim2)`
# for calls like `randexp(dim1, dim2)`
return :($f($SM{$(escall(fargs)...)}))
elseif length(fargs) == 3
# for calls like `rand(rng, dim1, dim2)`
# for calls like `rand(type, dim1, dim2)`
# for calls like `rand(sampler, dim1, dim2)`
# for calls like `randn(rng, dim1, dim2)`
# for calls like `randn(type, dim1, dim2)`
# for calls like `randexp(rng, dim1, dim2)`
# for calls like `randexp(type, dim1, dim2)`
return :($_f_with_Val($SM, $(esc(fargs[1])), Val($(esc(fargs[2]))), Val($(esc(fargs[3])))))
elseif length(fargs) == 4
# for calls like `rand(rng, type, dim1, dim2)`
# for calls like `rand(rng, sampler, dim1, dim2)`
# for calls like `randn(rng, type, dim1, dim2)`
# for calls like `randexp(rng, type, dim1, dim2)`
return :($_f_with_Val($SM, $(esc(fargs[1])), $(esc(fargs[2])), Val($(esc(fargs[3]))), Val($(esc(fargs[4])))))
else
error("@$SM got bad expression: $(ex)")
end
else
error("@$SM only supports the zeros(), ones(), rand(), randn(), and randexp() functions.")
error("@$SM only supports the zeros(), ones(), fill(), rand(), randn(), and randexp() functions.")
end
else
error("Bad input for @$SM")
Expand Down
Loading
Loading