From 86ef978ecb18c31fa92de10a26c8dee1a74b5631 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Fri, 24 Jul 2026 12:13:27 -0400 Subject: [PATCH] ## Summary Adds `ones_map` and `fill_map` alongside the existing `zeros_map`/`randn_map`/`rand_map` split constructors, plus the flat `ones`/`fill` axis-friendly companions they forward through. They follow the same shape as the rest of the family: `ones_map([T,] codomain_axes, domain_axes)` fills a map-shaped array with ones and `fill_map(v, codomain_axes, domain_axes)` with a constant value, flattening to `(codomain_axes..., conj.(domain_axes)...)` and defaulting the element type from `T`/`v`. This rounds out the value-filling `*_map` set so a backend can build a map-shaped (or charged) tensor with any of the standard fills, not only zeros/randn/rand. --- Project.toml | 2 +- src/similar_map.jl | 44 ++++++++++++++++++++++++++++++---------- test/test_similar_map.jl | 22 +++++++++++++++++++- 3 files changed, 55 insertions(+), 13 deletions(-) diff --git a/Project.toml b/Project.toml index fcad8e03..1ce25441 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorAlgebra" uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -version = "0.17.9" +version = "0.17.10" authors = ["ITensor developers and contributors"] [workspace] diff --git a/src/similar_map.jl b/src/similar_map.jl index 25b417d7..4cd22f54 100644 --- a/src/similar_map.jl +++ b/src/similar_map.jl @@ -33,23 +33,31 @@ end """ zeros([T,] axes) -> A + ones([T,] axes) -> A randn([rng,] [T,] axes) -> A rand([rng,] [T,] axes) -> A + fill(v, 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). +Axis-friendly counterparts of `Base.zeros`/`Base.ones`/`Base.randn`/`Base.rand`/`Base.fill`, +taking the axes as a single tuple. `Base.zeros`/`Base.ones`/`Base.fill` already accept 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 ones end function randn end function rand end +function fill end +@doc (@doc zeros) ones @doc (@doc zeros) randn @doc (@doc zeros) rand +@doc (@doc zeros) fill zeros(::Type{T}, axes::Tuple) where {T} = Base.zeros(T, axes) +ones(::Type{T}, axes::Tuple) where {T} = Base.ones(T, axes) +fill(value, axes::Tuple) = Base.fill(value, 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) @@ -63,29 +71,43 @@ end """ zeros_map([T,] codomain_axes, domain_axes) -> M + ones_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 + fill_map(v, 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 +filled with zeros (`zeros_map`), ones (`ones_map`), normally-distributed values +(`randn_map`), uniformly-distributed values (`rand_map`), or the value `v` (`fill_map`), +with element type `T` (defaulting to `Float64`; `fill_map` takes it from `v`). 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)/[`ones`](@ref)/[`randn`](@ref)/[`rand`](@ref)/[`fill`](@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 ones_map end function randn_map end function rand_map end +function fill_map end +@doc (@doc zeros_map) ones_map @doc (@doc zeros_map) randn_map @doc (@doc zeros_map) rand_map +@doc (@doc zeros_map) fill_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 +ones_map(codomain_axes, domain_axes) = ones_map(Float64, codomain_axes, domain_axes) +function ones_map(::Type{T}, codomain_axes, domain_axes) where {T} + return ones(T, (codomain_axes..., conj.(domain_axes)...)) +end +function fill_map(value, codomain_axes, domain_axes) + return fill(value, (codomain_axes..., conj.(domain_axes)...)) +end for f in (:randn_map, :rand_map) g = Symbol(chopsuffix(String(f), "_map")) diff --git a/test/test_similar_map.jl b/test/test_similar_map.jl index 772f848f..7446b6e2 100644 --- a/test/test_similar_map.jl +++ b/test/test_similar_map.jl @@ -1,5 +1,6 @@ using Random: default_rng -using TensorAlgebra: TensorAlgebra, rand_map, randn_map, similar_map, zeros_map +using TensorAlgebra: + TensorAlgebra, fill_map, ones_map, rand_map, randn_map, similar_map, zeros_map using Test: @test, @testset @testset "similar_map ($T)" for T in (Float32, Float64, ComplexF32, ComplexF64) @@ -30,6 +31,14 @@ end @test z isa Matrix{T} @test size(z) == (2, 3) @test iszero(z) + o = TensorAlgebra.ones(T, ax) + @test o isa Matrix{T} + @test size(o) == (2, 3) + @test all(isone, o) + fl = TensorAlgebra.fill(T(2), ax) + @test fl isa Matrix{T} + @test size(fl) == (2, 3) + @test all(==(T(2)), fl) for f in (TensorAlgebra.randn, TensorAlgebra.rand) a = f(default_rng(), T, ax) @test a isa Matrix{T} @@ -48,6 +57,17 @@ end @test size(z) == (2, 3, 4) @test iszero(z) + o = ones_map(T, cod, dom) + @test o isa Array{T, 3} + @test size(o) == (2, 3, 4) + @test all(isone, o) + @test ones_map(cod, dom) isa Array{Float64, 3} # eltype defaults to Float64 + + fl = fill_map(T(3), cod, dom) + @test fl isa Array{T, 3} + @test size(fl) == (2, 3, 4) + @test all(==(T(3)), fl) + for f in (randn_map, rand_map) # Fully specified. a = f(default_rng(), T, cod, dom)