tf.keras.models.Model使用時のエラー: ValueError: The two structures don't have the same nested structure

Question

前提・実現したいこと

tensorflow/kerasでLayerの自作に取り組んでいる者です

ﾊﾞｯﾁ数x高さx幅xﾁｬﾝﾈﾙ数で構成される入力に対して、サイズを一回り縮小かつ1のみで
構成されたテンソルを返すlayerを作成しているのですがtf.keras.models.ModelにLayerを
組み込んだ際に

ValueError: The two structures don't have the same nested structure.

とエラーが出てきてしまいます。
おそらく、Layer内部のtf.while_loopの設定が間違っていると思うのですが、アドバイスいただけたら幸いです

よろしくお願いいたします。

該当のソースコード

import numpy as np
import tensorflow as tf
from tensorflow.random import set_seed, uniform
from tensorflow.keras.layers import Layer, MaxPool2D, Input

set_seed(10)

#---- サイズを一回り縮小した、1のみで構成されたテンソルを返すlayerを作成
class maxpool_indices(Layer):
    def __init__(self, pool_size=(2, 2), **kwargs):
        super().__init__()
        self.pool_size   = pool_size

    def call(self, inputs, *args, **kwargs):

        #----inputの形からoutputの形を計算
        kmax = tf.shape(inputs)[0]
        imax = tf.shape(inputs)[1]-self.pool_size[0]+1
        jmax = tf.shape(inputs)[2]-self.pool_size[1]+1
        lmax = tf.shape(inputs)[3]

        kijl_max = kmax*imax*jmax*lmax

        #----繰り返し計算で1を積み上げていく -> [1,1,......,1]
        def body(kijl, max_indices):

            max_indices_new = tf.constant([1])
            max_indices  = tf.cond(kijl == 0, 
                                   lambda: max_indices_new,
                                   lambda: tf.concat([max_indices, max_indices_new], axis=-1))
            return tf.add(kijl,1), max_indices

        def cond(kijl, max_indices):
            return tf.less(kijl, kijl_max)


        tmp, max_indices = tf.while_loop(cond,
                                         body,
                                         (tf.constant(0, name='init'), tf.constant([-1], name='dummy_const')),
                                         shape_invariants=tf.TensorShape([1, None]))
        #----[1,1,......,1]を ﾊﾞｯﾁ数x高さx幅xﾁｬﾝﾈﾙ数 に整形
        return tf.reshape(max_indices, (kmax, imax, jmax, lmax))

#----動作確認
batch = 1
width = 4
height = 4
channel = 1
input = tf.random.uniform(shape=[batch, height, width, channel],
                          minval=0, maxval=10,
                          dtype=tf.dtypes.int32 )

x = maxpool_indices()(input)
display(x)
'''
上記は動作確認済み（下記出力）
<tf.Tensor: shape=(1, 3, 3, 1), dtype=int32, numpy=
array([[[[1],
         [1],
         [1]],

        [[1],
         [1],
         [1]],

        [[1],
         [1],
         [1]]]], dtype=int32)>
'''

def custom_model():
    inputs = Input(shape=(height, width, channel))
    x = maxpool_indices()(inputs) ##########ここでエラーが検出される#########
    return tf.keras.models.Model(inputs,x)

cm = custom_model()

発生している問題・エラーメッセージ

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-5-aa34b0004977> in <module>()
     58     return tf.keras.models.Model(inputs,x)
     59 
---> 60 cm = custom_model()

5 frames
<ipython-input-5-aa34b0004977> in custom_model()
     55 def custom_model():
     56     inputs = Input(shape=(height, width, channel))
---> 57     x = maxpool_indices()(inputs)
     58     return tf.keras.models.Model(inputs,x)
     59 

/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
    968     if _in_functional_construction_mode(self, inputs, args, kwargs, input_list):
    969       return self._functional_construction_call(inputs, args, kwargs,
--> 970                                                 input_list)
    971 
    972     # Maintains info about the `Layer.call` stack.

/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer.py in _functional_construction_call(self, inputs, args, kwargs, input_list)
   1106       # Check input assumptions set after layer building, e.g. input shape.
   1107       outputs = self._keras_tensor_symbolic_call(
-> 1108           inputs, input_masks, args, kwargs)
   1109 
   1110       if outputs is None:

/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer.py in _keras_tensor_symbolic_call(self, inputs, input_masks, args, kwargs)
    838       return nest.map_structure(keras_tensor.KerasTensor, output_signature)
    839     else:
--> 840       return self._infer_output_signature(inputs, args, kwargs, input_masks)
    841 
    842   def _infer_output_signature(self, inputs, args, kwargs, input_masks):

/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer.py in _infer_output_signature(self, inputs, args, kwargs, input_masks)
    878           self._maybe_build(inputs)
    879           inputs = self._maybe_cast_inputs(inputs)
--> 880           outputs = call_fn(inputs, *args, **kwargs)
    881 
    882         self._handle_activity_regularization(inputs, outputs)

/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/impl/api.py in wrapper(*args, **kwargs)
    693       except Exception as e:  # pylint:disable=broad-except
    694         if hasattr(e, 'ag_error_metadata'):
--> 695           raise e.ag_error_metadata.to_exception(e)
    696         else:
    697           raise

ValueError: in user code:

    <ipython-input-5-aa34b0004977>:45 call  *
        tmp, max_indices = tf.while_loop(cond,
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/deprecation.py:602 new_func  **
        return func(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/control_flow_ops.py:2541 while_loop_v2
        return_same_structure=True)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/control_flow_ops.py:2738 while_loop
        back_prop=back_prop)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/while_v2.py:84 while_loop
        expand_composites=False)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/nest.py:533 assert_same_structure
        % (str(e), str1, str2))

    ValueError: The two structures don't have the same nested structure.
    
    First structure: type=tuple str=(<tf.Tensor 'maxpool_indices_6/init:0' shape=() dtype=int32>, <tf.Tensor 'maxpool_indices_6/dummy_const:0' shape=(1,) dtype=int32>)
    
    Second structure: type=TensorShape str=(1, None)
    
    More specifically: Substructure "type=tuple str=(<tf.Tensor 'maxpool_indices_6/init:0' shape=() dtype=int32>, <tf.Tensor 'maxpool_indices_6/dummy_const:0' shape=(1,) dtype=int32>)" is a sequence, while substructure "type=TensorShape str=(1, None)" is not
    Entire first structure:
    (., .)
    Entire second structure:
    .

試したこと

ここに問題に対して試したことを記載してください。

補足情報（FW/ツールのバージョンなど）

ここにより詳細な情報を記載してください。

Answer 1

下記への変更で解決しました
お手数おかけしました。

shape_invariants=(tf.TensorShape(1), tf.TensorShape([None])