Skip to content

Commit

Permalink
Implement n-otimes (#213)
Browse files Browse the repository at this point in the history
  • Loading branch information
KeitaNakamura authored Sep 4, 2024
1 parent 0840a34 commit 83751ec
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 0 deletions.
14 changes: 14 additions & 0 deletions src/ops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,20 @@ julia> A = x ⊗ y
@inline otimes(x1::AbstractTensor, x2::AbstractTensor, others...) = otimes(otimes(x1, x2), others...)
@inline otimes(x::AbstractTensor) = x

struct OTimes{n} end
otimes(n::Int) = OTimes{n}()

@inline Base.:^(x::AbstractVec, ::OTimes{1}) = x
@generated function Base.:^(x::AbstractVec{dim}, ::OTimes{n}) where {dim, n}
indices = Expr(:tuple, [Symbol(i) for i in 1:n]...)
expr = Expr(:call, :*, [:(x[$i]) for i in indices.args]...)
TT = Tensor{Tuple{Symmetry{Tuple{fill(dim,n)...}}}}
quote
@_inline_meta
@einsum $TT $indices -> $expr
end
end

"""
dot(x::AbstractTensor, y::AbstractTensor)
x ⋅ y
Expand Down
4 changes: 4 additions & 0 deletions test/ops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,10 @@ end
@test (@inferred normalize(z))::typeof(z) Array(z) / norm(Array(z))
@test (@inferred (x, y, x, y))::Tensor{Tuple{3,3,3,3}, T} x y x y
@test (@inferred (x))::typeof(x) x
@test (@inferred x^⊗(1))::Tensor{Tuple{3}, T} x
@test (@inferred x^⊗(2))::Tensor{Tuple{@Symmetry{3,3}}, T} x x
@test (@inferred x^⊗(3))::Tensor{Tuple{@Symmetry{3,3,3}}, T} x x x
@test (@inferred x^⊗(4))::Tensor{Tuple{@Symmetry{3,3,3,3}}, T} x x x x
# nonsquare
x = rand(Vec{3, T})
y = rand(Vec{2, T})
Expand Down

0 comments on commit 83751ec

Please sign in to comment.