code: def test_saving(self):
      # Torch GPU: `model.predict(ds.batch(4))` fails on device placement
      # JAX GPU: out[0] and ref_out don't match. May be concat feature order?
      cls = feature_space.FeatureSpace
      fs = feature_space.FeatureSpace(
          features={
              "float_1": cls.float(),
              "float_2": cls.float_normalized(),
              "float_3": cls.float_discretized(num_bins=3),
              "int_1": cls.integer_categorical(
                  max_tokens=5, num_oov_indices=2
              ),
              "int_2": cls.integer_hashed(num_bins=32),
              "int_3": cls.integer_categorical(max_tokens=5),
          },
          crosses=[
              cls.cross(("float_3", "int_1"), crossing_dim=32),
              cls.cross(("int_1", "int_2"), crossing_dim=32),
          ],
          output_mode="concat",
      )
      fs.adapt(
          self._get_train_data_dict(
              as_dataset=True, include_strings=False
          )
      )
      data = {
          key: value[0]
          for key, value in self._get_train_data_dict(
              include_strings=False
          ).items()
      }
      ref_out = fs(data)
  
      temp_filepath = os.path.join(self.get_temp_dir(), "fs.keras")
      fs.save(temp_filepath)
      fs = saving_api.load_model(temp_filepath)
  
      # Save again immediately after loading to test idempotency
      temp_filepath = os.path.join(self.get_temp_dir(), "fs2.keras")
      fs.save(temp_filepath)
  
      # Test correctness of the first saved FS
      out = fs(data)
      self.assertAllClose(out, ref_out)
  
      inputs = fs.get_inputs()
      outputs = fs.get_encoded_features()
      model = models.Model(inputs=inputs, outputs=outputs)
      ds = self._get_train_data_dict(
          as_dataset=True, include_strings=False
      )
      out = model.predict(ds.batch(4))
      self.assertAllClose(out[0], ref_out)
  
      # Test correctness of the re-saved FS
      fs = saving_api.load_model(temp_filepath)
      out = fs(data)
      self.assertAllClose(out, ref_out)
  
imports: import os
  import pytest
  from tensorflow import data
  from keras import backend
  from keras import layers
  from keras import models
  from keras import ops
  from keras import testing
  from keras.layers.preprocessing import feature_space
  from keras.saving import saving_api
label_desc: Use image node_id 13 for calling the FeatureSpaceTest obj's underlying member method code with example usage: obj.test_saving() without return types