Git Product home page Git Product logo

einsum.jl's Introduction

Einsum.jl

Einstein summation notation similar to numpy's einsum function (but more flexible!).

PackageEvaluator Package Build Package Status
Einsum Build Status License
Einsum Project Status: Inactive - The project has reached a stable, usable state but is no longer being actively developed; support/maintenance will be provided as time allows. - help wanted!

To install: Pkg.add("Einsum").

Documentation

Basics

If the destination array is preallocated, then use =:

A = zeros(5,6,7) # need to preallocate destination
X = randn(5,2)
Y = randn(6,2)
Z = randn(7,2)
@einsum A[i,j,k] = X[i,r]*Y[j,r]*Z[k,r]

If destination is not preallocated, then use := to automatically create a new array A with appropriate dimensions:

using Einsum
X = randn(5,2)
Y = randn(6,2)
Z = randn(7,2)
@einsum A[i,j,k] := X[i,r]*Y[j,r]*Z[k,r]

What happens under the hood

To see exactly what is generated, use @macroexpand:

@macroexpand @einsum A[i,j,k] = X[i,r]*Y[j,r]*Z[k,r]

The @einsum macro automatically generates code that looks much like the following (note that we "sum out" over the index r, since it only occurs on the right hand side of the equation):

for k = 1:size(A,3)
    for j = 1:size(A,2)
        for i = 1:size(A,1)
            s = 0
            for r = 1:size(X,2)
                s += X[i,r] * Y[j,r] * Z[k,r]
            end
            A[i,j,k] = s
        end
    end
end

In reality, this code will be preceded by the the neccessary bounds checks and allocations, and take care to use the right types and keep hygenic.

You can also use updating assignment operators for preallocated arrays. E.g., @einsum A[i,j,k] *= X[i,r]*Y[j,r]*Z[k,r] will produce something like

for k = 1:size(A,3)
    for j = 1:size(A,2)
        for i = 1:size(A,1)
            s = 0
            for r = 1:size(X,2)
                s += X[i,r] * Y[j,r] * Z[k,r]
            end
            A[i,j,k] *= s
        end
    end
end

@einsimd

This is a variant of @einsum which will put @simd in front of the innermost loop; e.g., @einsum A[i,j,k] = X[i,r]*Y[j,r]*Z[k,r] will result approximately in

for k = 1:size(A,3)
    for j = 1:size(A,2)
        for i = 1:size(A,1)
            s = 0
            @simd for r = 1:size(X,2)
                s += X[i,r] * Y[j,r] * Z[k,r]
            end
            A[i,j,k] = s
        end
    end
end

Whether this is a good idea or not you have to decide and benchmark for yourself in every specific case. @simd makes sense for certain kinds of operations; make yourself familiar with its documentation and the inner workings of it in general.

Other functionality

In principle, the @einsum macro can flexibly implement function calls within the nested for loop structure. For example, consider transposing a block matrix:

z = Any[rand(2,2) for i=1:2, j=1:2]
@einsum t[i,j] := transpose(z[j,i])

This produces a for loop structure with a transpose function call in the middle. Approximately:

for j = 1:size(z,1)
    for i = 1:size(z,2)
        t[i,j] = transpose(z[j,i])
    end
end

This will work as long the function calls are outside the array names. Again, you can use @macroexpand to see the exact code that is generated.

Related Packages:

  • TensorOperations.jl has less flexible syntax (and does not allow certain contractions), but can produce much more efficient code. Instead of generating “naive” loops, it transforms the expressions into optimized contraction functions and takes care to use a good (cache-friendly) order for the looping.

einsum.jl's People

Contributors

ahwillia avatar dfdx avatar ntezak avatar phipsgabler avatar tkelman avatar

Watchers

 avatar  avatar

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. 📊📈🎉

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google ❤️ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.