Repository: Theano/Theano

Issue:
Opt: merge switch(cond, a, b) + switch(cond, c, d)
While playing with ELUs, I found that the gradient is a bit silly:

``` python
import theano
T = theano.tensor
x = T.matrix('x')
# define forward expression
elu = T.switch(x > 0, x, T.exp(x) - 1)
# define backward expression
delta = T.matrix('delta')
grelu = theano.grad(None, x, known_grads={elu: delta})
# compile backward expression
fgrelu = theano.function([x, delta], grelu)
# print graph
from theano.printing import debugprint, pp
debugprint(fgrelu)
```

```
HostFromGpu [id A] ''   4
 |GpuElemwise{Composite{(Switch(i0, i1, i2) + Switch(i0, i2, (exp(i3) * i1)))}}[(0, 0)] [id B] ''   3
   |GpuElemwise{Composite{Cast{float32}(GT(i0, i1))},no_inplace} [id C] ''   2
   | |GpuFromHost [id D] ''   0
   | | |x [id E]
   | |CudaNdarrayConstant{[[ 0.]]} [id F]
   |GpuFromHost [id G] ''   1
   | |delta [id H]
   |CudaNdarrayConstant{[[ 0.]]} [id F]
   |GpuFromHost [id D] ''   0
```

Written a bit more concisely, it is:

```
switch(x > 0, delta, 0) + switch(x > 0, 0, exp(x) * delta)
```

Note that the switches share the same condition.

The following optimizer merges it nicely:

``` python
local_merge_added_switch = theano.gof.PatternSub(
    in_pattern=(T.add, (T.switch, 'x', 'a1', 'b1'), (T.switch, 'x', 'a2', 'b2')),
    out_pattern=(T.switch, 'x', (T.add, 'a1', 'a2'), (T.add, 'b1', 'b2')),
    name='Merge addition of switch with same cond.', allow_multiple_clients=True)
T.opt.register_canonicalize(local_merge_added_switch, 'fast_run',
    name='Merge addition of switch with same cond.')
```

The gradient now reduces to:

```
HostFromGpu [id A] ''   4
 |GpuElemwise{Composite{Switch(i0, i1, (exp(i2) * i1))}}[(0, 0)] [id B] ''   3
   |GpuElemwise{Composite{Cast{float32}(GT(i0, i1))},no_inplace} [id C] ''   2
   | |GpuFromHost [id D] ''   0
   | | |x [id E]
   | |CudaNdarrayConstant{[[ 0.]]} [id F]
   |GpuFromHost [id G] ''   1
   | |delta [id H]
   |GpuFromHost [id D] ''   0
```

i.e.

```
switch(x > 0, delta, exp(x) * delta)
```

Should this be added to Theano? If so, where?



Discussion:
It make sense to put it in Theano I think.

I would put this in theano/tensor/opt.py. Tests would be go
theano/tensor/tests/test_opt.py

Thanks!

On Fri, Dec 4, 2015 at 2:30 PM, Jan Schlüter notifications@github.com
wrote:

> While playing with ELUs, I found that the gradient is a bit silly:
> 
> import theano
> T = theano.tensor
> x = T.matrix('x')# define forward expression
> elu = T.switch(x > 0, x, T.exp(x) - 1)# define backward expression
> delta = T.matrix('delta')
> grelu = theano.grad(None, x, known_grads={elu: delta})# compile backward expression
> fgrelu = theano.function([x, delta], grelu)# print graphfrom theano.printing import debugprint, pp
> debugprint(fgrelu)
> 
> HostFromGpu [id A] ''   4
>  |GpuElemwise{Composite{(Switch(i0, i1, i2) + Switch(i0, i2, (exp(i3) \* i1)))}}[(0, 0)] [id B] ''   3
>    |GpuElemwise{Composite{Cast{float32}(GT(i0, i1))},no_inplace} [id C] ''   2
>    | |GpuFromHost [id D] ''   0
>    | | |x [id E]
>    | |CudaNdarrayConstant{[[ 0.]]} [id F]
>    |GpuFromHost [id G] ''   1
>    | |delta [id H]
>    |CudaNdarrayConstant{[[ 0.]]} [id F]
>    |GpuFromHost [id D] ''   0
> 
> Written a bit more concisely, it is:
> 
> switch(x > 0, delta, 0) + switch(x > 0, 0, exp(x) \* delta)
> 
> Note that the switches share the same condition.
> 
> The following optimizer merges it nicely:
> 
> local_merge_added_switch = theano.gof.PatternSub(
>     in_pattern=(T.add, (T.switch, 'x', 'a1', 'b1'), (T.switch, 'x', 'a2', 'b2')),
>     out_pattern=(T.switch, 'x', (T.add, 'a1', 'a2'), (T.add, 'b1', 'b2')),
>     name='Merge addition of switch with same cond.', allow_multiple_clients=True)
> T.opt.register_canonicalize(local_merge_added_switch, 'fast_run',
>     name='Merge addition of switch with same cond.')
> 
> The gradient now reduces to:
> 
> HostFromGpu [id A] ''   4
>  |GpuElemwise{Composite{Switch(i0, i1, (exp(i2) \* i1))}}[(0, 0)] [id B] ''   3
>    |GpuElemwise{Composite{Cast{float32}(GT(i0, i1))},no_inplace} [id C] ''   2
>    | |GpuFromHost [id D] ''   0
>    | | |x [id E]
>    | |CudaNdarrayConstant{[[ 0.]]} [id F]
>    |GpuFromHost [id G] ''   1
>    | |delta [id H]
>    |GpuFromHost [id D] ''   0
> 
> i.e.
> 
> switch(x > 0, delta, exp(x) \* delta)
> 
> Should this be added to Theano? If so, where?
> 
> —
> Reply to this email directly or view it on GitHub
> https://github.com/Theano/Theano/issues/3742.

I agree we should put it in Theano.
For the future, I'm wondering if we could generalize it to more elementwise operations. It should work for multiplication, subtraction, division, as well as add and mul with more than 2 inputs.


Resolving patch:
diff --git a/theano/tensor/opt.py b/theano/tensor/opt.py
--- a/theano/tensor/opt.py
+++ b/theano/tensor/opt.py
@@ -3806,6 +3806,32 @@ def local_div_switch_sink(node):
     return False
 
 
+# Merge add/sub/mul/div/minimum/maximum/... of switches sharing the same
+# condition, to enable further simplification of their branches
+# Example: switch(c, a, b) + switch(c, x, y) -> switch(c, a+x, b+y)
+@register_canonicalize
+@gof.local_optimizer([T.Elemwise])
+def local_merge_switch_same_cond(node):
+    scal = theano.scalar
+    # node must be binary elemwise or add or mul
+    if not isinstance(node.op, T.Elemwise) or not isinstance(
+            node.op.scalar_op, (scal.BinaryScalarOp, scal.Add, scal.Mul)):
+        return
+    # all inputs must be switch
+    if not all(s.owner and isinstance(s.owner.op, T.Elemwise) and
+               isinstance(s.owner.op.scalar_op, scal.Switch)
+               for s in node.inputs):
+        return
+    # all switch conditions must be the same
+    cond = node.inputs[0].owner.inputs[0]
+    if not all(s.owner.inputs[0] is cond for s in node.inputs[1:]):
+        return
+    # pull out switch
+    return [T.switch(cond,
+                     node.op(*[s.owner.inputs[1] for s in node.inputs]),
+                     node.op(*[s.owner.inputs[2] for s in node.inputs]))]
+
+
 #############
 # Tile Opts #
 #############


Test patch:
diff --git a/theano/tensor/tests/test_opt.py b/theano/tensor/tests/test_opt.py
--- a/theano/tensor/tests/test_opt.py
+++ b/theano/tensor/tests/test_opt.py
@@ -4914,6 +4914,35 @@ def test_broadcast3(self):
                     isinstance(node.op, theano.tensor.Elemwise)]) == 0
 
 
+class test_local_merge_switch_same_cond(unittest.TestCase):
+    def test_elemwise(self):
+        # float Ops
+        mats = theano.tensor.matrices('cabxy')
+        c, a, b, x, y = mats
+        s1 = T.switch(c, a, b)
+        s2 = T.switch(c, x, y)
+        for op in (T.add, T.sub, T.mul, T.true_div, T.int_div, T.floor_div,
+                   T.minimum, T.maximum, T.gt, T.lt, T.ge, T.le, T.eq, T.neq,
+                   T.pow):
+            g = optimize(FunctionGraph(mats, [op(s1, s2)]))
+            assert str(g).count('Switch') == 1
+        # integer Ops
+        mats = theano.tensor.imatrices('cabxy')
+        c, a, b, x, y = mats
+        s1 = T.switch(c, a, b)
+        s2 = T.switch(c, x, y)
+        for op in (T.and_, T.or_, T.xor,
+                   T.bitwise_and, T.bitwise_or, T.bitwise_xor):
+            g = optimize(FunctionGraph(mats, [op(s1, s2)]))
+            assert str(g).count('Switch') == 1
+        # add/mul with more than two inputs
+        u, v = theano.tensor.matrices('uv')
+        s3 = T.switch(c, u, v)
+        for op in (T.add, T.mul):
+            g = optimize(FunctionGraph(mats + [u, v], [op(s1, s2, s3)]))
+            assert str(g).count('Switch') == 1
+
+
 class T_local_sum_prod(unittest.TestCase):
     """
     Test sum/prod opts in opt.py