Added shape docs for AvgPool2D layer from Keras (Kotlin#211)

https://keras.io/api/layers/pooling_layers/average_pooling2d/
dosier · Sep 29, 2022 · 138a215 · 138a215
1 parent 6e57581
commit 138a215
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tensorflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/core/layer/pooling/AvgPool2D.kt b/tensorflow/src/main/kotlin/org/jetbrains/kotlinx/dl/api/core/layer/pooling/AvgPool2D.kt
@@ -14,7 +14,9 @@ import org.tensorflow.op.Ops
 /**
  * Average pooling layer for 2D inputs (e.g. images).
  *
- * NOTE: Works with tensors which must have rank 4 (batch, height, width, channels).
+ * __Input shape:__  4D tensor with shape `(batch_size, rows, cols, channels)`.
+ *
+ * __Output shape:__ 4D tensor with shape `(batch_size, pooled_rows, pooled_cols, channels)`.
  *
  * @property [poolSize] The size of the sliding window for each dimension of input tensor (pool batch, pool height, pool width, pool channels).
  * Usually, pool batch and pool channels are equal to 1.