Landing page example broken

[DeD1rk]

The example on https://www.tensorflow.org/ranking doesn't run succesfully out of the box.
I tried opening it on colab and running, and I get this:

Downloading and preparing dataset 1.15 GiB (download: 1.15 GiB, generated: 310.08 MiB, total: 1.45 GiB) to /root/tensorflow_datasets/mslr_web/10k_fold1/1.2.0...

Dl Completed...: 100%
1/1 [01:55<00:00, 46.88s/ url]
Dl Size...: 100%
1176/1176 [01:55<00:00, 26.24 MiB/s]
Extraction completed...: 100%
15/15 [01:55<00:00, 5.52s/ file]

Dataset mslr_web downloaded and prepared to /root/tensorflow_datasets/mslr_web/10k_fold1/1.2.0. Subsequent calls will reuse this data.
Epoch 1/3

/usr/local/lib/python3.10/dist-packages/keras/src/engine/functional.py:642: UserWarning: Input dict contained keys ['_mask'] which did not match any model input. They will be ignored by the model.
  inputs = self._flatten_to_reference_inputs(inputs)

---------------------------------------------------------------------------

ValueError                                Traceback (most recent call last)

[<ipython-input-3-d669266a645f>](https://localhost:8080/#) in <cell line: 41>()
     39     loss=tfr.keras.losses.SoftmaxLoss(),
     40     metrics=tfr.keras.metrics.get("ndcg", topn=5, name="NDCG@5"))
---> 41 model.fit(ds, epochs=3)

1 frames

[/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py](https://localhost:8080/#) in tf__train_function(iterator)
     13                 try:
     14                     do_return = True
---> 15                     retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
     16                 except:
     17                     do_return = False

ValueError: in user code:

    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py", line 1377, in train_function  *
        return step_function(self, iterator)
    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py", line 1360, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py", line 1349, in run_step  **
        outputs = model.train_step(data)
    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/training.py", line 1126, in train_step
        y_pred = self(x, training=True)
    File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/usr/local/lib/python3.10/dist-packages/keras/src/engine/input_spec.py", line 280, in assert_input_compatibility
        raise ValueError(

    ValueError: Exception encountered when calling layer 'model' (type Functional).
    
    Input 0 of layer "batch_normalization_1" is incompatible with the layer: expected axis 2 of input shape to have value 1, but received input with shape (None, None, 136)
    
    Call arguments received by layer 'model' (type Functional):
      • inputs={'_mask': 'tf.Tensor(shape=(None, None), dtype=bool)', 'doc_id': 'tf.Tensor(shape=(None, None), dtype=int64)', 'float_features': 'tf.Tensor(shape=(None, None, 136), dtype=float64)', 'query_id': 'tf.Tensor(shape=(None,), dtype=string)'}
      • training=True
      • mask=None

Here are the installed versions from pip freeze:

tensorflow==2.14.0
tensorflow-datasets==4.9.3
tensorflow-ranking==0.5.3
tfds-nightly==4.9.3.dev202310060044

[DeD1rk]

Update: it seems like it is because of changes to the dataset: the features are now grouped.

Using something like:

inputs = {
        "float_features": tf.keras.Input(
            shape=(None, 136), dtype=tf.float32, name="float_features"
        ),
    }

Seems to work.