import tensorflow as tf


def initial_layer_binned(
    initial_layer_features: tf.Tensor,
    cutoff_data: CutoffData,
    band: str,
    soft_onehot: Nonlinearity = Nonlinearity.SIGMOID
):
    batch_size, twice_window_size, channels = map(int, initial_layer_features.shape)
    nonlinearity = nonlinearity_fcn(soft_onehot)
    if channels == 3:
        scales = cutoff_data.dflux_dt_dflux_dtime_scales(band)
        cutoffs = cutoff_data.dflux_dt_dflux_dtime_cutoffs(band)

        cutoffs_batch_window = tf.expand_dims(tf.expand_dims(cutoffs, 0), 0)
        scales_batch_window = tf.expand_dims(
            tf.expand_dims(tf.expand_dims(scales, 0), 0), -1
        )
        init_layer_per_cutoff = tf.expand_dims(initial_layer_features, -1)
        graph_typecheck.assert_shape(
            cutoffs_batch_window, [1, 1, channels, cutoff_data.embedding_size]
        )
        graph_typecheck.assert_shape(scales_batch_window, [1, 1, channels, 1])
        graph_typecheck.assert_shape(
            init_layer_per_cutoff, [batch_size, twice_window_size, channels, 1]
        )
        result = nonlinearity(
            (init_layer_per_cutoff - cutoffs_batch_window) / scales_batch_window
        )
        return graph_typecheck.assert_shape(
            result, [batch_size, twice_window_size, channels, cutoff_data.embedding_size]
