diff --git a/Project.toml b/Project.toml index 9628e9a..e74ffa0 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorAlgebra" uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -version = "0.16.1" +version = "0.16.2" authors = ["ITensor developers and contributors"] [workspace] @@ -10,6 +10,7 @@ projects = ["benchmark", "dev", "docs", "examples", "test"] EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67" StridedViews = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143" TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6" @@ -29,6 +30,7 @@ EllipsisNotation = "1.8" LinearAlgebra = "1.10" MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6" Mooncake = "0.4.202, 0.5" +Random = "1.10" Strided = "2.6" StridedViews = "0.5" TensorKit = "0.17" diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 8e5c8a8..714c771 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -1,8 +1,9 @@ module TensorAlgebraTensorKitExt +using Random: AbstractRNG using TensorAlgebra: TensorAlgebra -using TensorKit: TensorKit, AbstractTensorMap, ProductSpace, numind, permute, space, - spacetype, zerovector!, ← +using TensorKit: TensorKit, AbstractTensorMap, ElementarySpace, ProductSpace, numind, + permute, space, spacetype, zerovector!, ← using TensorOperations: TensorOperations as TO # ============================ AbstractArray-vocabulary bridge ============================ @@ -26,6 +27,44 @@ function TensorAlgebra.similar_map( return similar(a, T, ProductSpace{S}(codomain_axes...), ProductSpace{S}(domain_axes...)) end +# =============================== zeros_map / randn_map / rand_map ======================== +# A `TensorMap` keeps its codomain and domain as separate `ProductSpace`s rather than a single +# flattened axis, so build the `codomain ← domain` space directly instead of the dense +# flatten-and-dualize fallback. As with `similar_map`, the axes arrive codomain-facing +# (un-dualized), which is TensorKit's own codomain/domain convention. The elementary space type +# `S` is passed to `_map_homspace` explicitly so the two dispatch entries per constructor can +# read it from whichever of the codomain/domain is non-empty and share one builder; an empty +# axis tuple gives the unit space `ProductSpace{S}()`. +function _map_homspace(::Type{S}, codomain_axes, domain_axes) where {S <: ElementarySpace} + return ProductSpace{S}(codomain_axes...) ← ProductSpace{S}(domain_axes...) +end +function TensorAlgebra.zeros_map( + ::Type{T}, codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}} + ) where {T, S <: ElementarySpace} + return TensorKit.zeros(T, _map_homspace(S, codomain_axes, domain_axes)) +end +function TensorAlgebra.zeros_map( + ::Type{T}, codomain_axes::Tuple{}, domain_axes::Tuple{S, Vararg{S}} + ) where {T, S <: ElementarySpace} + return TensorKit.zeros(T, _map_homspace(S, codomain_axes, domain_axes)) +end +for (f, g) in ((:randn_map, :randn), (:rand_map, :rand)) + @eval begin + function TensorAlgebra.$f( + rng::AbstractRNG, ::Type{T}, + codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}} + ) where {T, S <: ElementarySpace} + return TensorKit.$g(rng, T, _map_homspace(S, codomain_axes, domain_axes)) + end + function TensorAlgebra.$f( + rng::AbstractRNG, ::Type{T}, + codomain_axes::Tuple{}, domain_axes::Tuple{S, Vararg{S}} + ) where {T, S <: ElementarySpace} + return TensorKit.$g(rng, T, _map_homspace(S, codomain_axes, domain_axes)) + end + end +end + # ================================ bipermutedimsopadd! ===================================== # `dest = β * dest + α * permutedims(op.(src), (perm_codomain, perm_domain))`. Delegate to # TensorKit's TensorOperations interface: `tensoradd!` realizes the permutation, the `op === conj` @@ -84,4 +123,28 @@ function TensorAlgebra.default_contract_algorithm( return TensorAlgebra.ContractAlgorithm(TO.DefaultBackend()) end +# ================================== linear-combination broadcast ========================= +# A `TensorMap` is not an `AbstractArray`, so it needs a `BroadcastStyle` to broadcast lazily +# (otherwise Base tries to `collect` it). A linear combination flattens (via `tryflattenlinear`) +# to a `LinearBroadcasted` that materializes through `add!`/`bipermutedimsopadd!` above; the +# `copyto!` here is not piracy because `LinearBroadcasted` is TensorAlgebra-owned. Element-wise +# (nonlinear) broadcast is not a meaningful operation on a symmetric tensor, so it errors rather +# than dense-converting. +struct TensorMapStyle <: Base.Broadcast.BroadcastStyle end +Base.Broadcast.BroadcastStyle(::Type{<:AbstractTensorMap}) = TensorMapStyle() +Base.Broadcast.BroadcastStyle(s::TensorMapStyle, ::TensorMapStyle) = s +Base.Broadcast.BroadcastStyle(s::TensorMapStyle, ::Base.Broadcast.BroadcastStyle) = s +Base.Broadcast.broadcastable(a::AbstractTensorMap) = a + +function Base.copyto!(dest::AbstractTensorMap, src::TensorAlgebra.LinearBroadcasted) + return TensorAlgebra.add!(dest, src, true, false) +end + +function Base.copy(::Base.Broadcast.Broadcasted{TensorMapStyle}) + return error( + "element-wise broadcast is not supported for a `TensorMap`; only linear combinations \ + such as `a .+ b` and `2 .* a` are supported" + ) +end + end diff --git a/src/similar_map.jl b/src/similar_map.jl index d527564..25b417d 100644 --- a/src/similar_map.jl +++ b/src/similar_map.jl @@ -1,3 +1,5 @@ +using Random: Random, AbstractRNG + """ similar_map(prototype, [T,] codomain_axes, domain_axes) -> M @@ -28,3 +30,76 @@ end function similar_map(prototype, codomain_axes, domain_axes) return similar_map(prototype, eltype(prototype), codomain_axes, domain_axes) end + +""" + zeros([T,] axes) -> A + randn([rng,] [T,] axes) -> A + rand([rng,] [T,] axes) -> A + +Axis-friendly counterparts of `Base.zeros`/`Base.randn`/`Base.rand`, taking the axes +as a single tuple. `Base.zeros` already accepts axes, but `Base.randn`/`Base.rand` +accept only integer dims, so these fill that gap for dense `Base.OneTo` axes and +otherwise forward to `Base` (so a graded-axis backend that extends `Base.randn`/`rand` +on its axis type is picked up). These are the flat (non-map) companions of +[`zeros_map`](@ref). +""" +function zeros end +function randn end +function rand end +@doc (@doc zeros) randn +@doc (@doc zeros) rand + +zeros(::Type{T}, axes::Tuple) where {T} = Base.zeros(T, axes) +for (f, g) in ((:randn, :randn), (:rand, :rand)) + @eval begin + $f(rng::AbstractRNG, ::Type{T}, axes::Tuple) where {T} = Base.$g(rng, T, axes) + function $f( + rng::AbstractRNG, ::Type{T}, axes::Tuple{Base.OneTo, Vararg{Base.OneTo}} + ) where {T} + return Base.$g(rng, T, map(length, axes)) + end + end +end + +""" + zeros_map([T,] codomain_axes, domain_axes) -> M + randn_map([rng,] [T,] codomain_axes, domain_axes) -> M + rand_map([rng,] [T,] codomain_axes, domain_axes) -> M + +Construct an array shaped as a linear map from `domain_axes` to `codomain_axes`, +filled with zeros (`zeros_map`), normally-distributed values (`randn_map`), or +uniformly-distributed values (`rand_map`), with element type `T` (defaulting to +`Float64`). These are the value-filling companions of [`similar_map`](@ref): the +domain axes are given un-dualized (codomain facing) and stored dual, so the default +flattens to the axis-friendly [`zeros`](@ref)/[`randn`](@ref)/[`rand`](@ref) over +`(codomain_axes..., conj.(domain_axes)...)` (`conj` dualizes a graded axis and is a +no-op on a dense one). Backends with map-shaped storage (e.g. a `TensorMap`) overload +these to build the codomain/domain directly. +""" +function zeros_map end +function randn_map end +function rand_map end +@doc (@doc zeros_map) randn_map +@doc (@doc zeros_map) rand_map + +zeros_map(codomain_axes, domain_axes) = zeros_map(Float64, codomain_axes, domain_axes) +function zeros_map(::Type{T}, codomain_axes, domain_axes) where {T} + return zeros(T, (codomain_axes..., conj.(domain_axes)...)) +end + +for f in (:randn_map, :rand_map) + g = Symbol(chopsuffix(String(f), "_map")) + @eval begin + $f(codomain_axes, domain_axes) = + $f(Random.default_rng(), codomain_axes, domain_axes) + function $f(rng::AbstractRNG, codomain_axes, domain_axes) + return $f(rng, Float64, codomain_axes, domain_axes) + end + function $f(::Type{T}, codomain_axes, domain_axes) where {T} + return $f(Random.default_rng(), T, codomain_axes, domain_axes) + end + function $f(rng::AbstractRNG, ::Type{T}, codomain_axes, domain_axes) where {T} + return $g(rng, T, (codomain_axes..., conj.(domain_axes)...)) + end + end +end diff --git a/test/test_similar_map.jl b/test/test_similar_map.jl index d517ef9..772f848 100644 --- a/test/test_similar_map.jl +++ b/test/test_similar_map.jl @@ -1,4 +1,5 @@ -using TensorAlgebra: similar_map +using Random: default_rng +using TensorAlgebra: TensorAlgebra, rand_map, randn_map, similar_map, zeros_map using Test: @test, @testset @testset "similar_map ($T)" for T in (Float32, Float64, ComplexF32, ComplexF64) @@ -21,3 +22,44 @@ using Test: @test, @testset @test eltype(O3) === ComplexF32 @test size(O3) == (2, 3) end + +# The flat axis-friendly constructors fill Base's gap: `randn`/`rand` reject `Base.OneTo`. +@testset "flat construction ($T)" for T in (Float32, Float64, ComplexF32, ComplexF64) + ax = (Base.OneTo(2), Base.OneTo(3)) + z = TensorAlgebra.zeros(T, ax) + @test z isa Matrix{T} + @test size(z) == (2, 3) + @test iszero(z) + for f in (TensorAlgebra.randn, TensorAlgebra.rand) + a = f(default_rng(), T, ax) + @test a isa Matrix{T} + @test size(a) == (2, 3) + end +end + +# The dense map constructors flatten `(codomain_axes..., conj.(domain_axes)...)`; `conj` is a +# no-op on a dense axis, so the shape is the concatenation of codomain and domain lengths. +@testset "map construction ($T)" for T in (Float32, Float64, ComplexF32, ComplexF64) + cod = (Base.OneTo(2), Base.OneTo(3)) + dom = (Base.OneTo(4),) + + z = zeros_map(T, cod, dom) + @test z isa Array{T, 3} + @test size(z) == (2, 3, 4) + @test iszero(z) + + for f in (randn_map, rand_map) + # Fully specified. + a = f(default_rng(), T, cod, dom) + @test a isa Array{T, 3} + @test size(a) == (2, 3, 4) + # Element type defaults to `Float64`. + b = f(cod, dom) + @test b isa Array{Float64, 3} + @test size(b) == (2, 3, 4) + end + + # An empty domain (all-codomain map) is how a plain tensor is constructed. + zc = zeros_map(T, cod, ()) + @test size(zc) == (2, 3) +end diff --git a/test/test_tensorkitext.jl b/test/test_tensorkitext.jl index a9799d0..d6715ee 100644 --- a/test/test_tensorkitext.jl +++ b/test/test_tensorkitext.jl @@ -1,6 +1,9 @@ +using Base.Broadcast: broadcasted using StableRNGs: StableRNG -using TensorAlgebra: contract, matricize, similar_map, unmatricize -using TensorKit: @tensor, Rep, SU₂, U₁, fuse, isomorphism, randn, space, ←, ⊗ +using TensorAlgebra: TensorAlgebra, contract, matricize, rand_map, randn_map, similar_map, + tryflattenlinear, unmatricize, zeros_map +using TensorKit: + @tensor, AbstractTensorMap, Rep, SU₂, U₁, fuse, isomorphism, randn, space, ←, ⊗ using Test: @test, @test_throws, @testset # A shared bond contracts when it sits in one operand's domain and the other's codomain, i.e. @@ -99,4 +102,49 @@ using Test: @test, @test_throws, @testset sm_domain = similar_map(t_domain, elt, (), (A1, A2)) @test space(sm_domain) == space(t_domain) end + + # The map constructors build a `TensorMap` from the codomain/domain spaces directly rather + # than flattening, mirroring the `similar_map` space convention (domain given un-dualized). + @testset "map construction on spaces" begin + A1 = Rep[U₁](0 => 2, 1 => 1) + A2 = Rep[U₁](0 => 1, 1 => 1) + B = Rep[U₁](0 => 1, -1 => 2) + + for f in (randn_map, rand_map) + t = f(rng, elt, (A1, A2), (B,)) + @test t isa AbstractTensorMap + @test space(t) == ((A1 ⊗ A2) ← B) + end + z = zeros_map(elt, (A1, A2), (B,)) + @test z isa AbstractTensorMap + @test space(z) == ((A1 ⊗ A2) ← B) + @test iszero(z) + + # An empty domain gives an all-codomain `TensorMap`, the plain-tensor case. + tc = randn_map(rng, elt, (A1, A2), ()) + @test space(tc) == ((A1 ⊗ A2) ← one(A1)) + + # An empty codomain is the mirror case: the space type comes from the domain. + td = randn_map(rng, elt, (), (A1, A2)) + @test space(td) == (one(A1) ← (A1 ⊗ A2)) + zd = zeros_map(elt, (), (A1, B)) + @test space(zd) == (one(A1) ← (A1 ⊗ B)) + end + + # A linear combination of `TensorMap`s flattens to a `LinearBroadcasted` that materializes + # into a `TensorMap` destination via `copyto!`; a nonlinear broadcast has no linear form. + @testset "linear-combination broadcast" begin + V = Rep[SU₂](0 => 1, 1 // 2 => 2) + a = randn(rng, elt, V ← V) + b = randn(rng, elt, V ← V) + + lb = tryflattenlinear(broadcasted(+, a, broadcasted(*, 2, b))) + dest = similar(a) + copyto!(dest, lb) + @test dest ≈ a + 2 * b + + # A nonlinear (element-wise) broadcast is not expressible as a `LinearBroadcasted`. + @test isnothing(tryflattenlinear(broadcasted(*, a, b))) + @test_throws ErrorException copy(broadcasted(*, a, b)) + end end