Specifying shape of normal sampling with partially empty keras tensor

Hi,
I’ve been migrating my code from keras 2 using tf to keras 3 based on torch.
For part of my model I need to sample a tensor of shape Z_mean, where Z_mean is a keras tensor of shape (None, 1024, 12).
When I exectute my code (and the minimal working code below) I get the following error:
Traceback (most recent call last):

  • File “/lhome/USER/.cache/JetBrains/RemoteDev/dist/8af836b5b190f_pycharm-professional-2024.1.1/plugins/python/helpers-pro/pydevd_asyncio/pydevd_asyncio_utils.py”, line 117, in _exec_async_code*
  • result = func()*
  • File “<‘input’>”, line 1, in *
  • File “/lhome/USER/PycharmProjects/PROJECT_NAME/venv/lib/python3.10/site-packages/keras/src/random/random.py”, line 27, in normal*
  • return backend.random.normal(*
  • File “/lhome/USER/PycharmProjects/PROJECT_NAME/venv/lib/python3.10/site-packages/keras/src/backend/torch/random.py”, line 37, in normal*
  • return torch.normal(*
    TypeError: normal(): argument ‘size’ must be tuple of ints, not tuple

Minimal working code:

Z_mean = keras.Input(shape=(1024, 12), batch_size=None)
eps = keras.random.normal(shape=ko.shape(Z_mean))

Sampling a partially None-shaped tensor worked fine in Keras 2, why am I doing wrong?

Since your question is Keras-specific, you might want to cross-post it into their discussion board.