[OpRefactor] Make decomposition a method of the main Operator

[glassnotes]

Context: Operator overhaul.

Description of the Change: The static decomposition has been moved to the Operator class, and a new method decompose has been added.

Benefits: Decompositions can now be obtained from specific instances of operations, and can use their stored parameters and wires. The move to Operator will allow the decomposition to be used as an arithmetic operation when the Operator class is overhauled.

Possible Drawbacks: Josh has found some issues with the decompose method in the cases where some of the non-trainable parameters are required for constructing the decomposition.

Related GitHub Issues: N/A

[glassnotes]

This is super ugly. Is there any way around this?

[mariaschuld]

I fear this is a very typical example for the weaknesses of the queuing system, and it does not only appear here, but is a general pattern! Every time we create an operation to call its method which creates other operations (i.e. to simplify an already created hamiltonian) this happens.

But the alternative, to keep all methods static, is not a scalable option I fear.

I hope we can solve this elsewhere.

[codecov]

Codecov Report

Merging #1873 (713d4ec) into master (91479c6) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #1873   +/-   ##
=======================================
  Coverage   98.98%   98.98%           
=======================================
  Files         219      219           
  Lines       16509    16513    +4     
=======================================
+ Hits        16341    16345    +4     
  Misses        168      168           

Impacted Files	Coverage Δ
pennylane/ops/qubit/attributes.py	100.00% <ø> (ø)
pennylane/_qubit_device.py	98.69% <100.00%> (ø)
pennylane/operation.py	98.88% <100.00%> (+<0.01%)	⬆️
pennylane/ops/qubit/matrix_ops.py	99.14% <100.00%> (ø)
pennylane/tape/reversible.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 91479c6...713d4ec. Read the comment docs.

[mariaschuld]

Fantastic, thanks for tackling this.

Does it feel right? I think it is much better to use the internal information in decomposition instead of arguments, and the only disadvantage is a structural problem of queuing, so I am inclined to ignore it for now and rather make a better queueing system.

[mariaschuld]

I like this a lot!

[mariaschuld]

I know, this pattern that the developer needs to know that parameters[0] is D is not optimal, but I can't think of a better option.

[mariaschuld]

On the other hand, if PennyLane feeds op.parameters into decomposition then PennyLane makes the assumption that D is parameters[0], which is probably a lot less safe!

[mariaschuld]

I fear this is a very typical example for the weaknesses of the queuing system, and it does not only appear here, but is a general pattern! Every time we create an operation to call its method which creates other operations (i.e. to simplify an already created hamiltonian) this happens.

But the alternative, to keep all methods static, is not a scalable option I fear.

I hope we can solve this elsewhere.

[josh146]

I must admit, I can't stop thinking about this 🤔

But the alternative, to keep all methods static, is not a scalable option I fear.

Would it make sense to keep decompositions as a static method (same for terms, and matrix), but simply enforce that it takes the same arguments as __init__? I feel like this might solve the problem.

E.g.,

class MyOp(Operation):
    def __init__(self, H, x, t, wires):
        super().__init__(H.coeffs, x, t, wires)
        self.hamiltonian = H

    @Operation.decomposition()
    def decomposition(self, H, x, t, wires):
        return [qml.ApproxTimeEvolution(H, t, wires), qml.RX(x, wires[0])]

where @Operation.decomposition() is just some 'registration' function.

So you could then do either:

>>> MyOp.decomposition(H, x, t, wires=[0, 1])

or

>>> MyOp(H, x, t, wires=[0, 1]).decomposition()

The ambiguity is also removed, since neither the user nor the base class makes any 'assumptions' about the signature of decomposition --- it always just matches the __init__.

[mariaschuld]

Is this still valid, since it's only relevant for tests and there one can do _decomposition?

[mariaschuld]

Perfect!

It may feel weird for now, but I up once we have all bits tidied up in a systematic manner the benefit will become obvious!

[mariaschuld]

Oh, only consider my suggestion for the changelog, I think that would need to be updated?

[glassnotes]

Oh, only consider my suggestion for the changelog, I think that would need to be updated?

My bad, forgot to update that! Thanks :)

[josh146]

I have one worry here, but I'm not sure of the solution. Given that we are planning on making it so that only trainable parameters are stored in self.parameters, and devs should instead store non-trainable parameters as static attributes... how will we pass non-trainable parameters here?

@mariaschuld, any ideas? I am drawing blanks.

For example, consider the following:

class MyOp(Operation):

    def __init__(self, H, t, boolean_arg=True, wires=None):
        # store trainable parameters
        super().__init__(H.coeffs, t, wires=None)

        # after calling __init__, trainable parameters
        # are now stored as self.parameters
        assert self.parameters == H.coeffs, t

        # store non-trainable parameters as an attribute
        self.boolean_arg = boolean_arg
        self.H = H

    @staticmethod
    def _decomposition(H, t, boolean_arg=True, wires=None):
        # the parent class will attempt to pass H.coeffs and t,
        # not H and boolean_arg!

[josh146]

The only solution I can think of, is something along the lines of what PyTorch does (https://pytorch.org/docs/stable/notes/extending.html):

# Inherit from Function
class LinearFunction(Function):

    # Note that both forward and backward are @staticmethods
    @staticmethod
    # bias is an optional argument
    def forward(ctx, input, weight, bias=None):
        ctx.save_for_backward(input, weight, bias)
        output = input.mm(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output

    # This function has only a single output, so it gets only one gradient
    @staticmethod
    def backward(ctx, grad_output):
        # This is a pattern that is very convenient - at the top of backward
        # unpack saved_tensors and initialize all gradients w.r.t. inputs to
        # None. Thanks to the fact that additional trailing Nones are
        # ignored, the return statement is simple even when the function has
        # optional inputs.
        input, weight, bias = ctx.saved_tensors
        grad_input = grad_weight = grad_bias = None

        # These needs_input_grad checks are optional and there only to
        # improve efficiency. If you want to make your code simpler, you can
        # skip them. Returning gradients for inputs that don't require it is
        # not an error.
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(weight)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input)
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0)

        return grad_input, grad_weight, grad_bias

When defining custom PyTorch operations, you create a class with static methods.

However, the first argument of each static method is always ctx. This is a 'context' object that

1. provides helper methods, and
2. allows the developers to store anything

[josh146]

If we take a similar approach ourselves, the operator class could look like this:

class MyOp(Operation):

    @staticmethod
    def create(ctx, H, t, boolean_arg=True, wires=None):
        # store trainable parameters
        ctx.save_parameters(H.coeffs, t)

        # store non-trainable parameters as an attribute
        self.ctx.boolean_arg = boolean_arg
        self.ctx.H = H

    @staticmethod
    def _decomposition(*parameters, wires=None, ctx=None):
        H.coeffs, t = parameters
        H = ctx.Hamiltonian
        boolean_arg = ctx.boolean_arg

[josh146]

Forgot to pick this when I left my review :(

[josh146]

Looks good @glassnotes!

[mariaschuld]

🚀

[mariaschuld]

🚀

[mariaschuld]

Nice and simple!

[mariaschuld]

I turned "approve" to "request changes" due to a minor codecov issue. I think since we now test decomposition instead of decompose there is no test where an int is fed as a wire. Maybe a test like this should help:

def test_decomposition_with_int_wire()
    op1 = qml.Hadamard.decomposition(1)
    op2 = qml.Hadamard.decomposition([1])
    assert op1 == op2

[mariaschuld]

Approval for if CI is happy now :)