The Softmax Brick Works Only For Matrices about blocks HOT 10 CLOSED

After giving it another thought I am inclined to just have a wrapper that adapts 2D-only blocks to work with more-D data, like the one @bartvm wrote ages ago. Otherwise we can end up with lots of boilerplate code dealing with dimensions.

from blocks.

Just throwing in an idea here, but how about sandwiching the softmax operation in reshapes?

I'm thinking of something like this:

from theano import tensor

def nd_softmax(a, axis=-1):
    swapped_a = a.swapaxes(axis, -1)
    matrix_a = swapped_a.reshape((swapped_a.shape[:-1].prod(), a.shape[-1]))
    softmax_swapped_a = tensor.nnet.softmax(matrix_a).reshape(swapped_a.shape)
    return softmax_swapped_a.swapaxes(axis, -1)

The downside of that approach is that the reshape operations may complicate the graph optimization and prevent some optimizations being applied (I've run into that issue in the past).

from blocks.

You could make the reshaping conditional though, no?

shape, ndim = a.shape, a.ndim
if ndim > 2:
    a = a.reshape((a.shape[:-1].prod(), a.shape[-1]))
softmax_a = tensor.nnet.softmax(a)
if ndim  > 2:
    softmax_a = softmax_a.reshape(shape)

This seems like code that would be generally useful though, so I personally had something in mind like:

class As2D(Brick):
    def __init__(self, application_method):
        self.application_method = application_method

    def apply(self, x):
        reshaped_x = x.reshape(...)  # As above
        y = self.application_method(reshaped_x)
        reshaped_y = y.reshape(...) 

That way you can do

y = As2D(Softmax().apply).apply(x)

from blocks.

I like that idea. Do you think there's a use case for applying softmax on another axis than the last one? If so, something like

class WithAxesSwapped(Brick):
    def __init__(self, application_method, axis_0, axis_1):
        self.application_method = application_method
        self.axis_0 = axis_0
        self.axis_1 = axis_1

    def apply(self, x):
        swapped_x = x.swapaxes(self.axis_0, self.axis_1)
        y = self.application_method(reshaped_x)
        swapped_y = y.swapaxes(self.axis_0, self.axis_1)

would let you do

y = WithAxesSwapped(As2D(Softmax().apply), axis, -1).apply(x)

from blocks.

Actually the attention brick from here could use that, normalization is done along the first axis there.

from blocks.

I'll write WithAxesSwapped and As2D and make a pull request out of that.

from blocks.

Where would it be most natural to put these two bricks? In blocks.bricks directly?

from blocks.

I would start a new file bricks/wrappers.py for that. @bartvm?

from blocks.

SGTM

from blocks.

Effectively closed by #404

from blocks.