Resubmit "Keras grouped convolutions"

[lgeiger]

This PR resubmits the changes from #36773 which were rolled back in dd2ea87.

I couldn't reproduce the failure mentioned in #36773 (comment) locally with or without XLA, so I am resubmitting the changes to see if CI is happy now.
@tanzhenyu could you take a look?

[lgeiger]

@tanzhenyu Any chance you could approve this so CI can run? Since this is a git revert of the commit that rolled back the merge of #36773, there should be no need for an indepth review again.

[lgeiger]

I rebased the PR and resolved the merge conflicts. @gbaned is there anything that still blocks this?

[tanzhenyu]

I'm not sure why it didn't show up in external CI. But here's what we got:

  File "/tensorflow/python/keras/layers/convolutional_test.py", line 296, in test_conv3d
    self._run_test(kwargs, expected_output_shape)
  File "/tensorflow/python/keras/layers/convolutional_test.py", line 265, in _run_test
    expected_output_shape=expected_output_shape)
  File "/tensorflow/python/framework/test_util.py", line 1717, in decorated
    result = f(self, *args, **kwargs)
  File "/tensorflow/python/keras/testing_utils.py", line 227, in layer_test
    model.train_on_batch(input_data, actual_output)
  File "/tensorflow/python/keras/engine/training.py", line 1476, in train_on_batch
    logs = train_function(iterator)
  File "/tensorflow/python/eager/def_function.py", line 766, in __call__
    result = self._call(*args, **kwds)
  File "/tensorflow/python/eager/def_function.py", line 826, in _call
    return self._stateless_fn(*args, **kwds)
  File "/tensorflow/python/eager/function.py", line 2812, in __call__
    return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
  File "/tensorflow/python/eager/function.py", line 1838, in _filtered_call
    cancellation_manager=cancellation_manager)
  File "/tensorflow/python/eager/function.py", line 1915, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/tensorflow/python/eager/function.py", line 549, in call
    ctx=ctx)
  File "/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.UnimplementedError:  Hit a case for convolution that is not implemented on GPU.
	 [[{{node cluster_131_1/xla_compile}}]] [Op:__inference_train_function_7135]

Either this doesn't pass the GPU test, or doesn't pass XLA test

[lgeiger]

@tanzhenyu Thanks for taking another look, very much appreciated.

After some digging I was able to run the tests with XLA on a cloud VM and could reproduce the failure locally. It looks like currently train_on_batch fails for 3D grouped convolutions when XLA is enabled.
Since the underlying failure is not related to this PR and would also appear when using grouped convolutions with tf.nn.conv3d on master, I disbled the relevant part of the test in 215161f. Please let me know if that works for now.

[lgeiger]

Thanks for approving, looks like CI is also happy this time 💚

[tanzhenyu]

Thanks for approving, looks like CI is also happy this time 💚

Great, thanks for fixing!