Add ones_map and fill_map split constructors

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 <support@itensor.org> and contributors"]
 
 [workspace]

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"))

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)