Fix channels_first error in 3D Convolution Layers

[Frightera]

Related to the issue #1384. In 3D Convolutions if the data format is specified as channels_first it throws this error:

TypeError: Failed to convert object of type <class 'list'> to Tensor. Contents: [None, 64, 1024, 32]. Consider casting elements to a supported type. Full traceback can be found on the related issue. This PR should fix that error.

Please find the gist here.

[SiegeLordEx]

Is this something that can be reproduced in a test (

probability/tensorflow_probability/python/layers/conv_variational_test.py

Line 180 in 0b71262

class ConvVariational(object):

)? Technically we already test this layer w/ channels_first inside Sequential (

probability/tensorflow_probability/python/layers/conv_variational_test.py

Line 595 in 0b71262

def _testLayerInSequential(self, layer_class): # pylint: disable=invalid-name

).

You can see how to run tests here: https://github.com/tensorflow/probability/blob/main/CONTRIBUTING.md#unit-tests

[Frightera]

@SiegeLordEx The test cases' input shape includes the batch size. Hence not reproducible. This error pops out when you try to define input_shape without including the batch_size, so it does not occur when input is specified with batch size.. Please find the demo here.

I've tested my changing also, it can also be found here.

[SiegeLordEx]

I'm asking for the reproduction to be checked into the tests.

[Frightera]

Yes, this can be reproduced in unit tests. Changing the following lines will reproduce:

https://github.com/Frightera/probability/blob/788e7b8c07010414c073ebb4e737632c2d72e786/tensorflow_probability/python/layers/conv_variational_test.py#L611-L614

Please find the gist here.

Edit: @SiegeLordEx If that was not what you've asked, can you be more specific? Thanks.

[Frightera]

@SiegeLordEx This was something can be reproduced in tests. Please see my comment above. The first commit should fix that error, second one is also for testing purposes. This was not reproducible in tests before since input_shape was not passed as parameter to layer_class.

What will be the state of this PR? Is there anything should I do now?

Thanks.

[SiegeLordEx]

Thanks for the fix!