Fix export_lib.make_tensor_spec

[Grvzard]

Fixes #19913

a simple repro from the original issue with minor changes:

import keras
import numpy as np

class SomeModel(keras.Model):
    def __init__(self, output_channels: int):
        super().__init__()
        self.output_channels = output_channels
        self.fc1 = keras.layers.Dense(256, activation="relu")
        self.fc2 = keras.layers.Dense(
            self.output_channels * 2,
            bias_initializer="ones",
            kernel_initializer="zeros",
        )

    def build(self, input_shape):
        super().build(input_shape)
        self.fc1.build(input_shape)
        fc1_out = self.fc1.compute_output_shape(input_shape)
        self.fc2.build(fc1_out)

    def call(self, inputs):
        features = self.fc1(inputs)
        return self.fc2(features)


class SomeParentModel(keras.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.fx = SomeModel(20)

    def call(self, im_batch):
        features = self.fx(im_batch)
        ...
        return features

parent_inst = SomeParentModel()
random_data = np.random.random((1, 512)).astype('float32')

parent_inst(random_data)
parent_inst.save('test_model.keras')

reloaded = keras.models.load_model('test_model.keras', custom_objects={'SomeParentModel': SomeParentModel})
reloaded(random_data)
reloaded.export('test')

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 83.33333% with 1 line in your changes missing coverage. Please review.

Project coverage is 79.02%. Comparing base (f0bae91) to head (aca0141).
Report is 2 commits behind head on master.

Files	Patch %	Lines
keras/src/export/export_lib.py	83.33%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #19915   +/-   ##
=======================================
  Coverage   79.02%   79.02%           
=======================================
  Files         499      499           
  Lines       46433    46436    +3     
  Branches     8545     8548    +3     
=======================================
+ Hits        36692    36695    +3     
  Misses       8015     8015           
  Partials     1726     1726           

Flag	Coverage Δ
keras	78.88% <83.33%> (+<0.01%)	⬆️
keras-jax	62.42% <83.33%> (+<0.01%)	⬆️
keras-numpy	57.22% <0.00%> (-0.01%)	⬇️
keras-tensorflow	63.64% <83.33%> (+<0.01%)	⬆️
keras-torch	62.38% <0.00%> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[fchollet]

Thanks for the fix! Can you add a unit test for this case?

[fchollet]

LGTM, thank you!