yhtang / funfact Goto Github PK

For constant optimizable offsets in a tensor expression, we should add the support for zero-dimensional tensors.

An example use case and suggested syntax is given in the following RBF expression:

# RBF approximation, one instance
m, n = 24, 36
r = 8
u = ff.tensor('u', m, r)
v = ff.tensor('v', n, r)
a = ff.tensor('a', r)
b = ff.tensor('b')
i, j, k = ff.indices('i, j, k')
tsrex = a[k] * ff.exp(-(u[i, ~k] - v[j, ~k])**2) + b[:] # or b[...]

Showcase quantum circuit synthesis using FF.

Outline:

• Introduction of quantum compilation/quantum circuit synthesis
• Toy problem description
• Visualization of hypothesis circuit topology
• Model setup
• Optimization
• Result analysis

Describe the bug
Accessing elements or slices of a vectorized factorization model fails.

To Reproduce

A = ff.tensor('A', 5, 2)
B = ff.tensor('B', 2, 4)
i, j, k = ff.indices('i, j, k')
tsrex = A[i, j] * B[j, k]
model = Factorization(tsrex)
model[0, 0]
model_vec = Factorization(tsrex, nvec=2)
model_vec[0, 0, 0]

The first one works, the second one errors: TypeError: indices() missing 1 required positional argument: 'slices'

Expected behavior
The elementwise access should return the requested element of the vectorized factorization.

docs/pages/user-guide/tensor-and-indices.md

This is the first part of the user guide. The goal is to:

• Introduce the APIs to create tensors and indices
• !!! note: explain the 'abstract' nature of the tensors, i.e. they are not yet populated with concrete numbers yet.
• !!! note: explain the index identity logic: identities check is performed on UUIDs on a Python variable basis, not done on the 'symbol', which is merely for human visualization.

Describe the bug
Since einop expects both operands to be tensors, a literal value of native Python type must be first promoted before usage.

To Reproduce

a = ff.tensor('a', 2, 3)
i, j = ff.indices(2)
fac = ff.Factorization.from_tsrex(a[i, j]**2)
fac()

Leads to

TypeError: Value LiteralValue(2, None) with type <class 'funfact.lang._terminal.LiteralValue'> is not a valid JAX type

Expected behavior
Literal values should be regarded as 0-th order tensors but without the need to be indexed.

Showcasing RBF approximation of dense matrices.

Outline:

• Introduction to low-complexity matrix approximation (motivation, SVD, NMF, etc.)
• RBF approximation
• Model setup
• Optimization
• Analysis

Certain properties of the AST can be of good practical value for users, yet these are only accessible through a long detour at the moment. For example, the set of 'live' indices, the shape of the resulting tensor, etc. It would be very helpful to expose them as @property of the TsrEx class.

Properties to expose:

• set of live indices
• shape
• ndim
• einspec
• ...

Describe the bug
Our current einop implementation depends on order='F' of numpy.transpose or jax.numpy.transpose, which PyTorch does not implement. Simply removing the order specifier leads to incorrect results.

To Reproduce

ff.use('torch')
u = ff.tensor(np.arange(3))
v = ff.tensor(np.ones(4))
i = ff.index()
fac = Factorization(u[[*i]] * v[[*i]])
fac()

Expected behavior
[0., 0., 0., 0., 1., 1., 1., 1., 2., 2., 2., 2.]

Additional context
This StackOverflow answer provides a solution that is potentially slow. Also, I tried to directly cast the solution and ended up getting a single number instead of 12 elements.

Describe the bug
The LeafInitializer interpreter has an optional argument dtype which defaults to ab.float32. However, it is not used by any of the primitives.

Expected behavior
The data type of the created tensor should be specified by this argument.

Additional context
Having ab.float32 as the default value in the function signature causes very early import of the default backend. This is an undesired behavior. Since the initializer is very likely an internal tool used only by developers, I suggest we remove the default values.

docs/pages/user-guide/tsrex.md

This is the second part of the user guide. Our goal here is to showcase how to express computations on tensors using either an indexless or Einstein-notation format.

!!! note
Emphasize that FunFact adopts a lazy evaluation model, in contrast to the eager mode as enabled by default by PyTorch and JAX. Computation does not happen immediately until at funfact.from_tsrex or funfact.factorize.

Outline:

• Indexless expressions (mention this is basically the same as NumPy): elementwise operations (+, -, *, /), inner product (@), exponentiation (**), etc.
• Indxed expressions: this invokes the Einstein-summation convention where repeated indices are contracted/reduced.
  • ~ for explicitly disabling reduction
  • * for within-dimension outer product (i.e. Kronecker)
  • >> for transposition
  • repeated application of [...] leads to index renaming.
  • !!! note: explain the difference between [...] >> [...] and [...][...].
• Query properties of the tensor expression: ndim, shape etc.

Describe the bug
#40 introduces the automatic application of several interpreters whenever a tensor expression is created. However, since the tensor expressions are constructed recursively, this means that the analyzers are also called repeatedly. This seems to significantly slow down the creation of tensor expressions.

To Reproduce
TODO: time the creation of tensor expressions. My impression is that creating a Tucker expression now could take 100s of milliseconds.

Related to #54, #57

Add full support for vectorization to the factorization model.

Some considerations:

• It probably makes sense to consider every factorization as vectorized even if nvec=1
• Bookkeeping of all the best solutions can happen in a copy
• Total loss = sum over vectorized loss function

This is a follow-up to #33. It seems like both the index propagator, slice propagator, and einop method needs to be updated accordingly. Meanwhile, given that there are now three types of indices (live, keep, and kron) in the einop operation, it might make sense to bypass the einspec procedure and directly use the live indices to carry out einop operations.

For penalty terms enforcing (unitary) constraints we need to be able to evaluate the Hermitian conjugate of a tensor expression. Transposition is already implemented, so we have to add a feature conj to evaluate the complex conjugate of a tensor expression.

Path optimization should only happen between commutative contractions with the same precedence. If connected non-commutative operations share the same user-override precedence, a warning should be raised.

Describe the bug
When a single index is used to index a tensor, the index node always has the Kronecker product flag set.

To Reproduce

x = ff.tensor('x', 5)
i = ff.index('i')
x[i].asciitree

leads to

 x[*i] 
 ├── x 
 ╰── *i 
     ╰── *i 

Expected behavior
No Kronecker product flag shall be marked on the index node.

 x[i] 
 ├── x 
 ╰── i 
     ╰── i 

A clear and concise description of what you expected to happen.

Analysis
This issue occurs become we implement the __iter__ method to enable the * syntax for indices. However, whenever a single index is used to address tensor, it becomes the sole argument to __getitem__, which in turn triggers the unpacking behavior (i.e. __iter__). However, when a single index is used to address tensor, the __getitem__ method of tensor expressions would attempt to test if it is iterable by checking for the existence of __iter__, thus triggering the setting of the kron flag.

For potential complications in setting up GPU versions of tensor backends, we can point users to the respective documentation page of JAX/PyTorch, etc.

Also, mention the different [] flavors that can be specified.

Describe the bug
ff.exp(2) leads to exp(None) in LaTeX rendering.

This is because the LiteralValue class does not properly fall back to str when the latex attribute is None.

The einsum API of common numerical linear algebra libraries such as NumPy, PyTorch, and JAX all support the extended Einstein notation format, which allows an explicit expression of the left-hand side indices in the specs string. This explicit specification is the part after the -> symbol:

# implicit lhs, inferred to be ik
ij,jk

# explicit lhs
ij,jk->ik

# explicit lhs, result transposed
ij,jk->ki

# explicit lhs, disables summation over j
ij,jk->ijk

As such, the specification of left-hand side indices has the effect of transposition and explicitly enabling/disabling the reduction operation over certain indices.

Adding this functionality would entail semantics that may look like:

# A, B, C, D: tensors
# i, j, k: indices
C[k, i] = A[i, j] * B[j, k]
D[i, k, j] = A[i, j] * B[j, k]

Showcase using FF to implement compressed layers in deep neural networks.

Outline:

• Introduction on NN compression/tensorization.
• Example problem description (MNIST)
• NN architecture visualization
• Model setup
• Training
• Analysis

docs/pages/user-guide/eval.md

This is the third part of the user guide. Our goal here is to introduce:

• How to instantiate factorization objects from tensor expressions.
• Query properties of the factorization: ndim, shape, etc. Say this inherits from the tensor expressions.
• Obtain elements of the final output:
  • fac() for getting the entire tensor
  • fac[] for element/fiber-wise access
  • fac[] for factor access
  • .factors and .all_factors for factor access

Describe the bug
While using the [...] syntax to rename indices, if the new indices overlap with the dummy indices (integrated out thus do not appear in the final live indices), then this overlap is not detected and duplicate indices will be rendered.

To Reproduce
Steps or a code snippet to reproduce the behavior:

u = ff.tensor('u', 2, 3)
v = ff.tensor('v', 3, 5)
i, j, k = ff.indices('i, j, k')
tsrex = u[i, j] * v[j, k]
tsrex[i, j]

gives

u[i, j] * v[j, j]

which results in incorrect Einstein notation.

Expected behavior
u[i, _0] * v[_0, j] where _0 is an anonymous dummy index created on-demand.

The current tensor expressions are lazily evaluated --- defining them only specifies the AST, and calculation only happens when we explicitly apply an evaluator. This is useful for training, but could be cumbersome for interferencing. Ideally, we might need 3 levels of eagerness:

• lazy (-1): composing tensor expressions only constructs AST
• default (0): perform certain static analysis such as index propagation, shape check, etc.
• eager (+1): immediate initialization and evaluation of the concrete tensors.

This could be implemented as a context manager following the no_grad/set_grad_on/... convention.

Following the lines of thought in #44, eager evaluation can be triggered using a @property such as value. The challenge here is to avoid repeated initialization of the leaves, and determine how much caching is possible when the leaf data is subject to modification.

Describe the bug
If the spec string generated for einop doesn't contain output indices, the method fails at the parsing step.

Steps to reproduce
The following code:

A = ff.tensor('A', 2, 3)
B = ff.tensor('B', 2, 3)
i, j, k = ff.indices('i, j, k')
tsrex = A[[i, j]] * B[[i, j]]
fac = ff.Factorization(tsrex)
fac()

returns ValueError: not enough values to unpack (expected 4, got 3)

Expected behavior
The code above should return a scalar value.

Add functionality to perform (partial) trace operations on repeated indices of a single index_notation node.

Examples:

A = tensor('A', 8, 8)
B = tensor('B', 2, 3, 4, 3)
i, j, k = indices('i, j, k')
traceA = A[i, i]
partial_traceB = B[i, j, k, j]

This could be implemented at the level of the index_notation node.

There are a few refactoring issues to resolve for the initial release.

Move all files to appropriate (new) modules:

• _einop should be moved from the lang>interpreters to algorithm module
• look into reorganizing the optim module
• nn neural network module
• rethink the interface to Factorization and the factorize function: 1) fit method for factorization, 2) add vectorize option, 3) tol convergence checking, ...
• ...

Expose all required functionality to the enduser: Factorization, factorize, neural network, ...

e.g. minplus, Viberti, etc.

We can currently evaluate a factorization model elementwise or slicewise, i.e.:

fac[0, 1, 2]
fac[1:2, :, 3]

are supported for factorization models. It would be useful to add support for Ellipsis indexing similar to numpy:

fac[0, ...]
fac[..., 1]
fac[1:2, ..., 3]

This is particularly useful for accessing instances of a vectorized model.

docs/pages/user-guide/factorize.md

Introduces the funfact.factorize method:

• Most basic, one-click usage.
• Vectorization for multi-instance learning
  • explain how to obtain results
• Fine-tuning: loss, integrator, tolerance, steps, etc.

Future: upon implementation of shape inference, explain the shape deduction logic and how to query for missing shapes.

Something more intuitive and interesting?

Currently all tensors in a tensor expression are updated in the gradient descent step of a factorization model. We should implement a feature such that the user can decide which of the tensors (and maybe also LiteralValues?) are actually optimized.

Potential syntax:

A = ff.tensor('A', 2, 3, optimizable=True)
offset = ff.literal('o', optimizable=True)
...

Currently, the einop method has to distinguish the access between broadcasting and non-broadcasting dimensions, resulting in a number of if-else code blocks. We can replace that with the boadcast_to method, which is supported by all currently implemented backends. The broadcast view will then transparently handle data access and eliminate the need for special treatment from our end.

Referfences:
https://numpy.org/doc/stable/reference/generated/numpy.broadcast_to.html
https://pytorch.org/docs/stable/generated/torch.broadcast_to.html
https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.broadcast_to.html

The actual relevant functions are numpy.expand_dims and torch.unsqueeze:

• https://numpy.org/doc/stable/reference/generated/numpy.expand_dims.html
• https://pytorch.org/docs/stable/generated/torch.unsqueeze.html

Implement a preliminary set of unit tests that checks the core functionalities such as tensor expression definition and evaluation.

Describe the bug
Potnetial flaw in the .T[...] transposition syntax:

When people write $A^T_{ij}$, they mean the (i, j)-th element of the transposition of A. Therefore, the Einstein notation $A_{ij} \cdot A^T_{ij}$ means the inner product between A and its transposition.

In Funfact, however, if A is a tsrex with live indices i, j, then A.T[i. j] does nothing --- it merely means to reorder the live indices of A from i, j to i, j, so nothing happens.

This is because .T means tranpose according to the following [] instead of directly creating a transposed tensor where the dimensions are reversed.

This could catch a user off-guard if they are trying to do a literal translation of some formula on paper to code using FunFact.

On the other hand, with the Einstein notation, there is no need to use .T to obtain the transposition of a tensor for computation --- simply adjusting the order of the indices will do. The only purpose of .T is to actually transpose the tensor in memory for storage purposes.

To Reproduce

i, j, k = ff.indices('i, j, k')
A = ff.tensor('A', 3, 3)
tsrex = A * A.T[j, k]  # expects to get $A \cdot A^T$, won't work
tsrex = A[i, j] * A[k, j]  # this works and computes A * A^T

Discussion

Is there a need to adjust this?

We currently have the Adam optimizer and the MSE (L2) loss function.

We should add a few more options to the optimization library:

Optimizers

• RMSProp
• AdaGrad
• ...

Loss functions

• L1 loss
• KL divergence
• ...

Additionally, we should allow the user to add a penalty term to the loss function to enforce (unitary) constraints.

Candidate syntax:

(A[i, k] * B[k, j]) @ (i, j, k)

(A[i, ~k] * B[~k, j])

Describe the bug
The shape analyzer doesn't correctly check the sizes of contracting indices.

To Reproduce

A = ff.tensor('A', 2, 3)
B = ff.tensor('B', 3, 4)
i, j, k = ff.indices('i, j, k')
tsrex = A[[i, j]] * B[[i, j]]
tsrex.shape

Expected behavior
The code above returns () while it is expected to raise a SyntaxError as the contracting indices don't have compatible dimensions.