Update group_normalization.py

[charbull]

opening a pull request to fix the len while saving the model as suggested by @freedomtan

reference context: #1033
specific issue: #1034

/usr/local/lib/python3.7/dist-packages/keras_cv/models/generative/stable_diffusion/internal/layers/group_normalization.py in _create_broadcast_shape(self, input_shape)
85
86 def _create_broadcast_shape(self, input_shape):
---> 87 broadcast_shape = [1] * len(input_shape)
88 broadcast_shape[self.axis] = input_shape[self.axis] // self.groups

TypeError: Exception encountered when calling layer 'group_normalization_60' (type GroupNormalization).

len is not well defined for a symbolic Tensor (Shape:0). Please call x.shape rather than len(x) for shape information.

Call arguments received by layer 'group_normalization_60' (type GroupNormalization):
• args=('tf.Tensor(shape=(None, 64, 64, 320), dtype=float32)',)
• kwargs=<class 'inspect._empty'>

What does this PR do?

changes from len(input) to input.shape.rank

Fixes # (issue)
#1034

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue? Please add a link
  to it if that's the case.
  The issue discussed with @LukeWood. This is just a minor patch.
• Did you write any new necessary tests?
• If this adds a new model, can you run a few training steps on TPU in Colab to ensure that no XLA incompatible OP are used?

[bhack]

@LukeWood @ianstenbit Why this was not already catched by our jit_compile=True XLA?

[ianstenbit]

Thanks @charbull!
This looks good to me -- it's consistent with how I implemented GroupNorm in core Keras (which is now available in 2.11).

I'd like to eventually update our SD implementation to use Keras GroupNorm instead of our own, but that'll add a 2.11 dependency which I don't think we want to do just yet.

I am going to manually test this change against some SD workflows (too large to run on GitHub tests). If they look good, I will merge this.

@bhack - I suspect we didn't encounter this because this seems to be TFLite specific?

[ianstenbit]

/gcbrun

[ianstenbit]

LGTM -- thank you!

My test workflows all look good with this change, and it seems like a low-risk improvement