keras/src/export/export_lib.py/make_tensor_spec
def _get_input_signature(model):
    shapes_dict = getattr(model, "_build_shapes_dict", None)
    if not shapes_dict:
        return None

    def make_tensor_spec(structure):
        # We need to turn wrapper structures like TrackingDict or _DictWrapper
        # into plain Python structures because they don't work with jax2tf/JAX.
        if isinstance(structure, dict):
            return {k: make_tensor_spec(v) for k, v in structure.items()}
        elif isinstance(structure, tuple):
            if all(isinstance(d, (int, type(None))) for d in structure):
                return tf.TensorSpec(
                    shape=(None,) + structure[1:], dtype=model.input_dtype
                )
            return tuple(make_tensor_spec(v) for v in structure)
        elif isinstance(structure, list):
            if all(isinstance(d, (int, type(None))) for d in structure):
                return tf.TensorSpec(
                    shape=[None] + structure[1:], dtype=model.input_dtype
                )
            return [make_tensor_spec(v) for v in structure]
        else:
            raise ValueError(
                f"Unsupported type {type(structure)} for {structure}"
            )

    return [make_tensor_spec(value) for value in shapes_dict.values()]