Merge pull request #33419 from Prakharprasun:doc-lax-conv-transpose

PiperOrigin-RevId: 839829550

Author: Google-ML-Automation

jax/_src/lax/convolution.py

@@ -302,6 +302,13 @@ def conv_transpose(lhs: Array, rhs: Array, strides: Sequence[int],
   This function directly calculates a fractionally strided conv rather than
   indirectly calculating the gradient (transpose) of a forward convolution.
 
+  Notes:
+    TensorFlow/Keras Compatibility: By default, JAX does NOT reverse the
+    kernel's spatial dimensions. This differs from TensorFlow's "Conv2DTranspose"
+    and similar frameworks, which flip spatial axes and swap input/output channels.
+
+    To match TensorFlow/Keras behavior, set "transpose_kernel=True" .
+
   Args:
     lhs: a rank `n+2` dimensional input array.
     rhs: a rank `n+2` dimensional array of kernel weights.