Repository: tensorflow
Issue: Custom Layer output shape isn't working with model.summary()
Labels: stat:awaiting response,stale,comp:keras,TF 2.18

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

2.18

### Custom code

Yes

### OS platform and distribution

Ubuntu 24.04 WSL

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

model.summary() output

**Current Behavior**
```
Model: "Q_KAN"
____________________________________________________________________________
 Layer (type)                Output Shape              Param #   Trainable  
============================================================================
 input_68 (InputLayer)       [(None, 2)]               0         Y          
                                                                            
 DenseKAN (DenseQKan)        None                      18        Y          
                                                                            
 RescalePi (Rescale)         None                      0         N          
                                                                            
============================================================================
Total params: 18 (144.00 Byte)
Trainable params: 18 (144.00 Byte)
Non-trainable params: 0 (0.00 Byte)
____________________________________________________________________________
```

**Expected Behavior**
```
Model: "Q_KAN"
____________________________________________________________________________
 Layer (type)                Output Shape              Param #   Trainable  
============================================================================
 input_68 (InputLayer)       [(None, 2)]               0         Y          
                                                                            
 DenseKAN (DenseQKan)         [(None, 10)]                      18        Y          
                                                                            
 RescalePi (Rescale)          [(None, 10)]                      0         N          
                                                                            
============================================================================
Total params: 18 (144.00 Byte)
Trainable params: 18 (144.00 Byte)
Non-trainable params: 0 (0.00 Byte)
____________________________________________________________________________
```

### Standalone code to reproduce the issue

**Custom layer code:**

```python
class Rescale(tf.keras.layers.Layer):
    def __init__(self,scale_limit=np.pi,**kwargs):
        super().__init__(trainable=False,**kwargs)
        self.scale_limit = scale_limit
    
    def call(self,inputs):
        s = tf.reduce_sum(inputs,axis=-1)
        return (inputs/s)*self.scale_limit

class DenseQKan(tf.keras.layers.Layer):
    def __init__(self,units:int,circuit:qml.QNode,layers:int,**kwargs):
        super().__init__(**kwargs)
        self.circuit = circuit
        self.qubits =  len(circuit.device.wires)
        self.units = units
        self.qbatches = None
        self.layers = layers
        
    def build(self,input_shape):
        if input_shape[-1]> self.qubits:
            self.qbatches = np.ceil(input_shape[-1]/self.qubits).astype(np.int32)
        else:
            self.qbatches = 1
        self.layer_weights = []
        for u in range(self.units):
            self.layer_weights.append(self.add_weight(shape=(self.qbatches,input_shape[-1]//self.qbatches,self.layers),
                                   initializer=tf.keras.initializers.RandomUniform(minval=-np.pi, maxval=np.pi, seed=None),
                                   trainable=True))
        self.built = True
        # W = np.random.uniform(low=-np.pi,high=np.pi,size=(self.units,self.qbatches,self.qubits,self.layers))
    # @tf.function(reduce_retracing=True)

    def compute_output_shape(self,input_shape):
        print("Build Input Shape",input_shape)
        return (input_shape[0],self.units)
        
    def call(self,inputs):
        assert self.qbatches != None 
        splits = tf.split(inputs,self.qbatches,-1) 
        out = []
        for u in range(self.units):
            unit_out = 0
            for qb in range(self.qbatches):
                qb_out = tf.reduce_sum(tf.stack(self.circuit(splits[qb],self.layer_weights[u][qb]),axis=-1),axis=-1)
                unit_out = unit_out+qb_out
            out.append(unit_out)
        out = tf.stack(out,axis=-1)
        return out
```

**Model definition code:**
```python
# def create_model(units,qubits,layers,circuit,input_shape=2):
inp = Input(shape=input_shape)
out = DenseQKan(units,circuit,layers,name="DenseKAN")(inp)
out = Rescale(name="RescalePi")(out)
model = Model(inputs=inp,outputs=out,name="Q_KAN")
model.summary(show_trainable=True)
```

### Relevant log output

_No response_