import keras
from keras import layers, ops
import numpy as np

class ProdDenseLayer(layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        
    def build(self, input_shape):
        self.input_prod = ops.prod(input_shape[1:])
        self.dense = layers.Dense(self.input_prod, activation='relu')
        self.dense.build(input_shape)
        
    def call(self, inputs):
        scale_factor = ops.prod(ops.shape(inputs)[1:])
        scaled_inputs = inputs * ops.cast(scale_factor, inputs.dtype)
        return self.dense(scaled_inputs)

# Main model using ops.prod
class ProdModel(keras.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.prod_layer1 = ProdDenseLayer()
        
    def build(self, input_shape):
        self.prod_layer1.build(input_shape)
        
    def call(self, inputs):
        batch_size = ops.shape(inputs)[0]
        total_elements = ops.prod(ops.shape(inputs)[1:])
        normalized_inputs = inputs / ops.cast(total_elements, inputs.dtype)
        x = self.prod_layer1(normalized_inputs)
        return x

# Create dummy data
batch_size = 32
input_shape = (10,)
X_train = np.random.randn(batch_size * 10, *input_shape).astype(np.float32)
y_train = np.random.randint(0, 2, (batch_size * 10, 1)).astype(np.float32)
model = ProdModel()
model.compile(
    optimizer='adam',
    loss='binary_crossentropy',
    metrics=['accuracy']
)
history = model.fit(
    X_train, y_train,
    epochs=3,
    batch_size=batch_size,
    validation_split=0.2,
    verbose=1
)
model.summary()