Address second-round review findings

Latent bugs fixed:
- EinsumDense bias ordering: Keras stores the bias variable with dims in rhs
  char order, not bias_axes-string order. Reorder bias_axes_ chars to match
  rhs_ before computing strides so the right element is read.
- EinsumDense output_shape index: Keras's output_shape excludes the batch
  dim, so indexing rhs[i] against full_output_shape_[i] was off by one for
  output-only chars. Prepend a -1 placeholder for the batch char in the
  C++ creator to keep positional indexing consistent.

Defensive guards:
- IntegerLookup now requires num_oov_indices == 1 (was <= 1, which silently
  accepted strict mode that Keras would have raised on).
- Add n_units > 0 assertion to lstm_impl, gru_impl, simple_rnn_impl.
- Add groups > 0 / group_volume > 0 assertions to GroupNormalization.
- Add out_size > 0 assertion to AdaptivePooling adapt_range.

Style:
- Make IntegerLookup's vocab_index_ const and populate via a helper called
  from the member init list, matching the rest of the new layer files.

Test coverage:
- LSTM had return_state=True but GRU and SimpleRNN didn't — added.
- GroupedQueryAttention with num_query_heads == num_key_value_heads now also
  exercises distinct K/V sequence lengths (cross-attention shape).
- IntegerLookup is now reached by an end-to-end model (chained after
  Discretization to get an integer-valued tensor).
- Several more Random* augmentation passthroughs (RandomContrast,
  RandomRotation, RandomTranslation, RandomZoom, RandomHue,
  RandomSaturation, RandomSharpness) now go through the test pipeline.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: Dobiasd

include/fdeep/import_model.hpp

@@ -825,7 +825,10 @@ namespace internal {
         const std::string bias_axes = config["bias_axes"].is_null()
             ? std::string("")
             : std::string(config["bias_axes"]);
+        // Keras's output_shape excludes the batch dimension. Prepend -1 for
+        // the batch char so the layer can index it positionally against rhs.
         std::vector<int> output_shape;
+        output_shape.push_back(-1);
         for (const auto& dim : config["output_shape"])
             output_shape.push_back(dim.is_null() ? -1 : static_cast<int>(dim));
 

include/fdeep/layers/adaptive_pooling_3d_layer.hpp

@@ -53,6 +53,7 @@ namespace internal {
         {
             if (in_size == 1)
                 return { 0, 1 };
+            assertion(out_size > 0, "AdaptivePooling output_size must be > 0.");
             const auto start = static_cast<std::size_t>(std::floor(
                 static_cast<double>(i * in_size) / static_cast<double>(out_size)));
             const auto end = static_cast<std::size_t>(std::ceil(

include/fdeep/layers/einsum_dense_layer.hpp

@@ -195,20 +195,29 @@ namespace internal {
         }
 
         // Adds bias broadcast over the configured bias_axes (a subset of rhs).
+        // Keras stores the bias variable with its dims in rhs char order, not
+        // in the literal bias_axes string order (e.g. bias_axes='ed' on
+        // rhs='abde' still produces a (d, e)-shaped bias). Reorder the chars
+        // before computing strides so we read the right element.
         void add_bias(float_vec& out,
             const std::map<char, std::size_t>& char_size,
             const std::vector<std::size_t>& rhs_strides) const
         {
             if (bias_axes_.empty())
                 return;
+            std::string bias_chars;
+            for (char c : rhs_)
+                if (bias_axes_.find(c) != std::string::npos)
+                    bias_chars.push_back(c);
+
             const auto& bias_vals = *bias_.as_vector();
-            const auto bias_sizes = sizes_for_chars(bias_axes_, char_size);
+            const auto bias_sizes = sizes_for_chars(bias_chars, char_size);
             const auto bias_strides = compute_strides(bias_sizes);
 
             std::map<char, std::size_t> pos;
             for (std::size_t out_idx = 0; out_idx < out.size(); ++out_idx) {
                 decode_index(out_idx, rhs_, rhs_strides, pos);
-                out[out_idx] += bias_vals[encode_offset(bias_axes_, bias_strides, pos)];
+                out[out_idx] += bias_vals[encode_offset(bias_chars, bias_strides, pos)];
             }
         }
 

include/fdeep/layers/group_normalization_layer.hpp

@@ -49,13 +49,15 @@ namespace internal {
                 "GroupNormalization is currently only supported on the last (channel) axis.");
 
             const std::size_t channels = in_shape.depth_;
+            assertion(groups_ > 0, "GroupNormalization groups must be > 0.");
             assertion(channels % groups_ == 0,
                 "Number of channels must be divisible by number of groups.");
             const std::size_t channels_per_group = channels / groups_;
 
             const std::size_t spatial_volume = in_shape.size_dim_5_
                 * in_shape.size_dim_4_ * in_shape.height_ * in_shape.width_;
             const std::size_t group_volume = spatial_volume * channels_per_group;
+            assertion(group_volume > 0, "GroupNormalization input has zero volume.");
 
             const auto& src = *input.as_vector();
             float_vec means(groups_, 0);

include/fdeep/layers/integer_lookup_layer.hpp

@@ -28,23 +28,31 @@ namespace internal {
             , mask_token_(mask_token)
             , num_oov_indices_(num_oov_indices)
             , vocab_offset_((has_mask_token ? 1 : 0) + num_oov_indices)
-            , vocab_index_()
+            , vocab_index_(build_vocab_index(vocabulary, num_oov_indices))
         {
-            // Keras hashes OOV inputs across multiple OOV buckets via FarmHash,
-            // which fdeep does not currently replicate. Restrict to a single
-            // OOV bucket (the common case) to keep behavior bit-identical.
-            assertion(num_oov_indices_ <= 1,
-                "IntegerLookup with num_oov_indices > 1 is not supported.");
-            for (std::size_t i = 0; i < vocabulary.size(); ++i)
-                vocab_index_[vocabulary[i]] = i;
         }
 
     protected:
         const bool has_mask_token_;
         const std::int64_t mask_token_;
         const std::size_t num_oov_indices_;
         const std::size_t vocab_offset_;
-        std::unordered_map<std::int64_t, std::size_t> vocab_index_;
+        const std::unordered_map<std::int64_t, std::size_t> vocab_index_;
+
+        // num_oov_indices > 1 needs Keras's FarmHash, which fdeep does not
+        // currently replicate. num_oov_indices == 0 ("strict" mode) raises on
+        // OOV, also unsupported. Require exactly one OOV bucket.
+        static std::unordered_map<std::int64_t, std::size_t> build_vocab_index(
+            const std::vector<std::int64_t>& vocabulary, std::size_t num_oov_indices)
+        {
+            assertion(num_oov_indices == 1,
+                "IntegerLookup requires num_oov_indices == 1.");
+            std::unordered_map<std::int64_t, std::size_t> idx;
+            idx.reserve(vocabulary.size());
+            for (std::size_t i = 0; i < vocabulary.size(); ++i)
+                idx[vocabulary[i]] = i;
+            return idx;
+        }
 
         tensors apply_impl(const tensors& inputs) const override
         {
@@ -61,8 +69,7 @@ namespace internal {
                 if (it != vocab_index_.end()) {
                     out[i] = static_cast<float_type>(vocab_offset_ + it->second);
                 } else {
-                    // num_oov_indices_ <= 1 (asserted in ctor): OOV maps to 0,
-                    // shifted by 1 if a mask token occupies index 0.
+                    // OOV maps to 0, shifted by 1 if a mask token occupies index 0.
                     out[i] = static_cast<float_type>(has_mask_token_ ? 1 : 0);
                 }
             }

include/fdeep/recurrent_ops.hpp

@@ -125,6 +125,7 @@ namespace internal {
         const std::string& activation,
         const std::string& recurrent_activation)
     {
+        assertion(n_units > 0, "LSTM units must be > 0.");
         const MappedRowMajorMatrixXf W = eigen_row_major_mat_from_shared_values(
             weights.size() / (n_units * 4), n_units * 4,
             weights.data());
@@ -205,6 +206,7 @@ namespace internal {
         const std::string& activation,
         const std::string& recurrent_activation)
     {
+        assertion(n_units > 0, "GRU units must be > 0.");
         const std::size_t n_timesteps = input.shape().width_;
         const std::size_t n_features = input.shape().depth_;
 
@@ -300,6 +302,7 @@ namespace internal {
         const float_vec& bias,
         const std::string& activation)
     {
+        assertion(n_units > 0, "SimpleRNN units must be > 0.");
         const std::size_t n_timesteps = input.shape().width_;
         const std::size_t n_features = input.shape().depth_;
 

keras_export/generate_test_models.py

@@ -42,6 +42,8 @@
 from keras.layers import RNN, LSTMCell, GRUCell, SimpleRNNCell, StackedRNNCells
 from keras.layers import Discretization, IntegerLookup, Masking
 from keras.layers import RandomBrightness, RandomFlip, RandomCrop
+from keras.layers import RandomContrast, RandomRotation, RandomTranslation, RandomZoom
+from keras.layers import RandomHue, RandomSaturation, RandomSharpness
 from keras.models import Model, load_model, Sequential
 
 __author__ = "Tobias Hermann"
@@ -546,6 +548,10 @@ def get_test_model_exhaustive() -> Model:
         num_key_value_heads=2, use_gate=True)(inputs[49], inputs[49], inputs[49]))
     outputs.append(GroupQueryAttention(head_dim=4, num_query_heads=4,
         num_key_value_heads=4)(inputs[49], inputs[49], inputs[49]))
+    # num_query_heads == num_key_value_heads with distinct K/V seq lengths
+    # exercises the MHA-equivalence path with cross-attention shapes.
+    outputs.append(GroupQueryAttention(head_dim=4, num_query_heads=4,
+        num_key_value_heads=4)(inputs[49], inputs[50], inputs[51]))
     outputs.append(GroupQueryAttention(head_dim=3, num_query_heads=4,
         num_key_value_heads=2)(inputs[49], inputs[50], inputs[51]))
 
@@ -587,10 +593,21 @@ def get_test_model_exhaustive() -> Model:
     # Discretization on a float input.
     outputs.append(Discretization(bin_boundaries=[-0.5, 0.0, 0.5, 1.0])(inputs[49]))
 
+    # IntegerLookup: scaled to int via Discretization so the input is float.
+    outputs.append(IntegerLookup(vocabulary=[0, 1, 2, 3])(
+        Discretization(bin_boundaries=[-1.0, 0.0, 1.0])(inputs[49])))
+
     # Training-only Random* augmentation layers (passed through at inference).
     outputs.append(RandomBrightness(factor=0.1)(inputs[22]))
     outputs.append(RandomFlip(mode='horizontal')(inputs[22]))
     outputs.append(RandomCrop(height=26, width=28)(inputs[22]))
+    outputs.append(RandomContrast(factor=0.1)(inputs[22]))
+    outputs.append(RandomRotation(factor=0.1)(inputs[22]))
+    outputs.append(RandomTranslation(0.1, 0.1)(inputs[22]))
+    outputs.append(RandomZoom(0.1)(inputs[22]))
+    outputs.append(RandomHue(factor=0.1, value_range=(0.0, 1.0))(inputs[22]))
+    outputs.append(RandomSaturation(factor=0.1, value_range=(0.0, 1.0))(inputs[22]))
+    outputs.append(RandomSharpness(factor=0.1, value_range=(0.0, 1.0))(inputs[22]))
 
     shared_conv = Conv2D(1, (1, 1),
                          padding='valid', name='shared_conv', activation='relu')
@@ -901,6 +918,9 @@ def get_test_model_recurrent() -> Model:
     rnn_seq = SimpleRNN(6, return_sequences=True)(inputs)
     rnn_last = SimpleRNN(5, activation='tanh')(rnn_seq)
     rnn_relu_no_bias = SimpleRNN(4, activation='relu', use_bias=False)(inputs)
+    rnn_state_out, rnn_state_h = SimpleRNN(3, return_state=True)(inputs)
+
+    gru_state_out, gru_state_h = GRU(4, return_state=True)(inputs)
 
     bidi_lstm_seq = Bidirectional(LSTM(4, return_sequences=True))(inputs)
     bidi_lstm_last_concat = Bidirectional(LSTM(3))(bidi_lstm_seq)
@@ -941,8 +961,8 @@ def get_test_model_recurrent() -> Model:
     outputs = [
         lstm_last, lstm_no_bias, lstm_relu, lstm_no_unit_forget,
         lstm_state_out, lstm_state_h, lstm_state_c,
-        gru_last, gru_no_bias,
-        rnn_last, rnn_relu_no_bias,
+        gru_last, gru_no_bias, gru_state_out, gru_state_h,
+        rnn_last, rnn_relu_no_bias, rnn_state_out, rnn_state_h,
         bidi_lstm_last_concat,
         bidi_lstm_sum,
         bidi_gru_mul,