Define AlgorithmSelection and BroadcatMapConversion submodules

Author: mtfishman

NDTensors/src/NDTensors.jl

@@ -19,13 +19,15 @@ using Strided
 using TimerOutputs
 using TupleTools
 
-# TODO: Define an `AlgorithmSelection` module
 # TODO: List types, macros, and functions being used.
-include("lib/algorithm.jl")
+include("lib/AlgorithmSelection/src/AlgorithmSelection.jl")
+using .AlgorithmSelection: AlgorithmSelection
 include("lib/BaseExtensions/src/BaseExtensions.jl")
 using .BaseExtensions: BaseExtensions
 include("lib/SetParameters/src/SetParameters.jl")
 using .SetParameters
+include("lib/BroadcastMapConversion/src/BroadcastMapConversion.jl")
+using .BroadcastMapConversion: BroadcastMapConversion
 include("lib/Unwrap/src/Unwrap.jl")
 using .Unwrap
 include("lib/RankFactorization/src/RankFactorization.jl")

NDTensors/src/blocksparse/contract.jl

@@ -1,3 +1,5 @@
+using .AlgorithmSelection: Algorithm, @Algorithm_str
+
 function contract(
   tensor1::BlockSparseTensor,
   labelstensor1,

NDTensors/src/lib/AlgorithmSelection/src/AlgorithmSelection.jl

@@ -0,0 +1,3 @@
+module AlgorithmSelection
+include("algorithm.jl")
+end

NDTensors/src/lib/AlgorithmSelection/src/algorithm.jl

@@ -0,0 +1,39 @@
+"""
+    Algorithm
+
+A type representing an algorithm backend for a function.
+
+For example, a function might have multiple backend algorithm
+implementations, which internally are selected with an `Algorithm` type.
+
+This allows users to extend functionality with a new algorithm but
+use the same interface.
+"""
+struct Algorithm{Alg,Kwargs<:NamedTuple}
+  kwargs::Kwargs
+end
+
+Algorithm{Alg}(kwargs::NamedTuple) where {Alg} = Algorithm{Alg,typeof(kwargs)}(kwargs)
+Algorithm{Alg}(; kwargs...) where {Alg} = Algorithm{Alg}(NamedTuple(kwargs))
+Algorithm(s; kwargs...) = Algorithm{Symbol(s)}(NamedTuple(kwargs))
+
+Algorithm(alg::Algorithm) = alg
+
+# TODO: Use `SetParameters`.
+algorithm_string(::Algorithm{Alg}) where {Alg} = string(Alg)
+
+function Base.show(io::IO, alg::Algorithm)
+  return print(io, "Algorithm type ", algorithm_string(alg), ", ", alg.kwargs)
+end
+Base.print(io::IO, alg::Algorithm) = print(io, algorithm_string(alg), ", ", alg.kwargs)
+
+"""
+    @Algorithm_str
+
+A convenience macro for writing [`Algorithm`](@ref) types, typically used when
+adding methods to a function that supports multiple algorithm
+backends.
+"""
+macro Algorithm_str(s)
+  return :(Algorithm{$(Expr(:quote, Symbol(s)))})
+end

NDTensors/src/lib/AlgorithmSelection/test/Project.toml

@@ -0,0 +1,3 @@
+[deps]
+NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

NDTensors/src/lib/AlgorithmSelection/test/runtests.jl

@@ -0,0 +1,12 @@
+@eval module $(gensym())
+using Test: @test, @testset
+using NDTensors.AlgorithmSelection: Algorithm, @Algorithm_str
+@testset "AlgorithmSelection" begin
+  @test Algorithm"alg"() isa Algorithm{:alg}
+  @test Algorithm("alg") isa Algorithm{:alg}
+  @test Algorithm(:alg) isa Algorithm{:alg}
+  alg = Algorithm"alg"(; x=2, y=3)
+  @test alg isa Algorithm{:alg}
+  @test alg.kwargs == (; x=2, y=3)
+end
+end

NDTensors/src/lib/BlockSparseArrays/src/BlockSparseArrays.jl

@@ -1,4 +1,5 @@
 module BlockSparseArrays
+using ..AlgorithmSelection: Algorithm, @Algorithm_str
 using BlockArrays:
   AbstractBlockArray,
   BlockArrays,

NDTensors/src/lib/BlockSparseArrays/src/LinearAlgebraExt/LinearAlgebraExt.jl

@@ -1,9 +1,9 @@
 module LinearAlgebraExt
+using ...AlgorithmSelection: Algorithm, @Algorithm_str
 using BlockArrays: BlockArrays, blockedrange, blocks
 using ..BlockSparseArrays: SparseArray, nonzero_keys # TODO: Move to `SparseArraysExtensions` module, rename `SparseArrayDOK`.
 using ..BlockSparseArrays: BlockSparseArrays, BlockSparseArray, nonzero_blockkeys
 using LinearAlgebra: LinearAlgebra, Hermitian, Transpose, I, eigen, qr
-using ...NDTensors: Algorithm, @Algorithm_str # TODO: Move to `AlgorithmSelector` module.
 using SparseArrays: SparseArrays, SparseMatrixCSC, spzeros, sparse
 
 # TODO: Move to `SparseArraysExtensions`.

NDTensors/src/lib/NamedDimsArrays/src/broadcastmapconversion.jl renamed to NDTensors/src/lib/BroadcastMapConversion/src/BroadcastMapConversion.jl

@@ -1,7 +1,7 @@
+module BroadcastMapConversion
 # Convert broadcast call to map call by capturing array arguments
 # with `map_args` and creating a map function with `map_function`.
 # Logic from https://github.com/Jutho/Strided.jl/blob/v2.0.4/src/broadcast.jl.
-# TODO: Move to a `BroadcastMapConversion` module.
 
 using Base.Broadcast: Broadcasted
 
@@ -45,3 +45,4 @@ function apply_tuple(t::Tuple, args)
   ttail, newargs = apply_tuple(Base.tail(t), newargs1)
   return (t1, ttail...), newargs
 end
+end

NDTensors/src/lib/BroadcastMapConversion/test/runtests.jl

@@ -0,0 +1,14 @@
+@eval module $(gensym())
+using Test: @test, @testset
+using NDTensors.BroadcastMapConversion: map_function, map_args
+@testset "BroadcastMapConversion" begin
+  using Base.Broadcast: Broadcasted
+  c = 2.2
+  a = randn(2, 3)
+  b = randn(2, 3)
+  bc = Broadcasted(*, (c, a))
+  @test copy(bc) ≈ c * a ≈ map(map_function(bc), map_args(bc)...)
+  bc = Broadcasted(+, (a, b))
+  @test copy(bc) ≈ a + b ≈ map(map_function(bc), map_args(bc)...)
+end
+end

NDTensors/src/lib/NamedDimsArrays/src/NamedDimsArrays.jl

@@ -1,9 +1,4 @@
 module NamedDimsArrays
-using ..BaseExtensions: BaseExtensions
-
-# TODO: Move to a `BroadcastMapConversion` module.
-include("broadcastmapconversion.jl")
-
 include("traits.jl")
 include("randname.jl")
 include("abstractnamedint.jl")

NDTensors/src/lib/NamedDimsArrays/src/abstractnameddimsarray.jl

@@ -1,3 +1,6 @@
+using ..BaseExtensions: BaseExtensions
+
+# Some of the interface is inspired by:
 # https://github.com/invenia/NamedDims.jl
 # https://github.com/mcabbott/NamedPlus.jl
 

NDTensors/src/lib/NamedDimsArrays/src/broadcast.jl

@@ -1,4 +1,5 @@
 using Base.Broadcast: BroadcastStyle, AbstractArrayStyle, DefaultArrayStyle, Broadcasted
+using ..BroadcastMapConversion: map_function, map_args
 
 struct NamedDimsArrayStyle{N} <: AbstractArrayStyle{N} end
 

NDTensors/src/lib/TensorAlgebra/src/TensorAlgebra.jl

@@ -1,6 +1,6 @@
 module TensorAlgebra
+using ..AlgorithmSelection: Algorithm, @Algorithm_str
 using LinearAlgebra: mul!
-using ..NDTensors: Algorithm, @Algorithm_str
 
 include("bipartitionedpermutation.jl")
 include("fusedims.jl")

NDTensors/src/lib/algorithm.jl

@@ -1,8 +1,10 @@
 @eval module $(gensym())
 using Test: @testset
 @testset "Test NDTensors lib $lib" for lib in [
+  "AlgorithmSelection",
   "BaseExtensions",
   "BlockSparseArrays",
+  "BroadcastMapConversion",
   "DiagonalArrays",
   "NamedDimsArrays",
   "SetParameters",

NDTensors/test/lib/BaseExtensions.jl

@@ -20,16 +20,6 @@ end
 @safetestset "NDTensors" begin
   @testset "$filename" for filename in [
     "lib/runtests.jl",
-    ## "lib/BaseExtensions.jl",
-    ## "lib/BlockSparseArrays.jl",
-    ## "lib/DiagonalArrays.jl",
-    ## "lib/NamedDimsArrays.jl",
-    ## "lib/SetParameters.jl",
-    ## "lib/SmallVectors.jl",
-    ## "lib/SortedSets.jl",
-    ## "lib/TagSets.jl",
-    ## "lib/TensorAlgebra.jl",
-    ## "lib/Unwrap.jl",
     "linearalgebra.jl",
     "dense.jl",
     "blocksparse.jl",