Repository: google/flax

Issue:
Unify behavior of strides arg of Conv and ConvTranspose
Currently, `Conv` accepts an int for `strides` which is then broadcast to the correct shape. However, `ConvTranspose` accepts only an optional Sequence, so no broadcasting is done. 
For consistency, I think the behavior should be equal, thus ConvTranspose should also accept integers for strides.

This works:
```python
from flax import linen as nn
module = nn.Conv(16, (4, 4), strides=2)
```

This doesn't:
```python
from flax import linen as nn
module = nn.ConvTranspose(16, (4, 4), strides=2)
```

Discussion:

Patch:
diff --git a/flax/experimental/nnx/nnx/tracers.py b/flax/experimental/nnx/nnx/tracers.py
index 88266823..f64e8660 100644
--- a/flax/experimental/nnx/nnx/tracers.py
+++ b/flax/experimental/nnx/nnx/tracers.py
@@ -14,6 +14,9 @@
 
 # Taken from flax/core/tracer.py 🏴‍☠️
 
+import contextlib
+import dataclasses
+import threading
 import typing as tp
 
 import jax
@@ -41,19 +44,70 @@ def current_jax_trace() -> MainTrace:
   return get_top_trace(())
 
 
+def get_all_traces(pytree: tp.Union[tp.Any, Tracer]) -> tp.Set[MainTrace]:
+  """Returns True if all tracers have the same main trace."""
+  if isinstance(pytree, Tracer):
+    return {pytree._trace.main}
+  else:
+    return {
+      trace._trace.main
+      for trace in jax.tree_util.tree_leaves(pytree)
+      if isinstance(trace, Tracer)
+    }
+
+
+def trace_level(main):
+  """Returns the level of the trace of -infinity if it is None."""
+  if main:
+    return main.level
+  return float('-inf')
+
+
+@dataclasses.dataclass
+class TraceContext(threading.local):
+  nnx_trace_stack: tp.List[MainTrace] = dataclasses.field(
+    default_factory=lambda: [current_jax_trace()]
+  )
+
+
+TRACE_CONTEXT = TraceContext()
+
+
+@contextlib.contextmanager
+def nnx_trace(trace: MainTrace):
+  TRACE_CONTEXT.nnx_trace_stack.append(trace)
+  try:
+    yield
+  finally:
+    TRACE_CONTEXT.nnx_trace_stack.pop()
+
+
+def current_nnx_trace() -> MainTrace:
+  return TRACE_CONTEXT.nnx_trace_stack[-1]
+
+
 class TraceState(reprlib.Representable):
-  __slots__ = ['_jax_trace']
+  __slots__ = ['_jax_trace', '_nnx_trace']
 
   def __init__(self):
     self._jax_trace = current_jax_trace()
+    self._nnx_trace = current_nnx_trace()
 
   @property
   def jax_trace(self):
     return self._jax_trace
 
+  @property
+  def nnx_trace(self):
+    return self._nnx_trace
+
   def is_valid(self) -> bool:
-    return self._jax_trace is current_jax_trace()
+    return (
+      self._jax_trace is current_jax_trace()
+      and self._nnx_trace is current_nnx_trace()
+    )
 
   def __nnx_repr__(self):
     yield reprlib.Object(f'{type(self).__name__}')
     yield reprlib.Attr('jax_trace', self._jax_trace)
+    yield reprlib.Attr('nnx_trace', self._nnx_trace)
diff --git a/flax/experimental/nnx/nnx/transforms.py b/flax/experimental/nnx/nnx/transforms.py
index aedbb2ad..7292eacf 100644
--- a/flax/experimental/nnx/nnx/transforms.py
+++ b/flax/experimental/nnx/nnx/transforms.py
@@ -42,6 +42,7 @@ from flax.experimental.nnx.nnx import (
   filterlib,
   rnglib,
   spmd,
+  tracers,
   variables,
 )
 from flax.experimental.nnx.nnx.module import GraphDef, Module, ModuleMeta
@@ -195,10 +196,12 @@ def get_jitted_fn(_module_type: type[M], f, options: JITOptions) -> JittedFn[M]:
     if isinstance(states, State):
       states = (states,)
 
-    if 'rngs' in kwargs:
-      kwargs['rngs'] = rnglib.Rngs(kwargs['rngs'])
-    module = graphdef.merge(*states)
-    out = f(module, *args, **kwargs)
+    nnx_trace = tracers.get_top_trace((args, kwargs))
+    with tracers.nnx_trace(nnx_trace):
+      if 'rngs' in kwargs:
+        kwargs['rngs'] = rnglib.Rngs(kwargs['rngs'])
+      module = graphdef.merge(*states)
+      out = f(module, *args, **kwargs)
 
     updates = module.split()
     out = (updates, out)
@@ -476,8 +479,9 @@ def grad_apply(options: GradOptions, f, module: Module, *args, **kwargs):
   def grad_fn(diff: State):
     nonlocal graphdef
 
-    module = graphdef.merge(diff, nondiff)
-    out = f(module, *args, **kwargs)
+    with tracers.nnx_trace(tracers.get_top_trace(diff)):
+      module = graphdef.merge(diff, nondiff)
+      out = f(module, *args, **kwargs)
 
     updates, graphdef = module.split()
     if options.has_aux:
@@ -843,41 +847,41 @@ def scan_apply(
 
   # transpose axes state
   scan_states = tuple(
-    jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, axis, 0), axes_state)
-    for axes_state, axis in zip(scan_states, options.variable_axes.values())
+      jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, axis, 0), axes_state)
+      for axes_state, axis in zip(scan_states, options.variable_axes.values())
   )
   # transpose axes arg
   scan_args = jax.tree_util.tree_map(
-    lambda axis, node: jax.tree_util.tree_map(
-      lambda x: jnp.moveaxis(x, axis, 0), node
-    )
-    if axis is not None
-    else None,
-    options.in_args_axes,
-    args,
-    is_leaf=lambda x: x is None,
+      lambda axis, node: jax.tree_util.tree_map(
+          lambda x: jnp.moveaxis(x, axis, 0), node
+      )
+      if axis is not None
+      else None,
+      options.in_args_axes,
+      args,
+      is_leaf=lambda x: x is None,
   )
   broadcast_args = jax.tree_util.tree_map(
-    lambda axis, node: None if axis is not None else node,
-    options.in_args_axes,
-    args,
-    is_leaf=lambda x: x is None,
+      lambda axis, node: None if axis is not None else node,
+      options.in_args_axes,
+      args,
+      is_leaf=lambda x: x is None,
   )
   scan_kwargs = jax.tree_util.tree_map(
-    lambda axis, node: jax.tree_util.tree_map(
-      lambda x: jnp.moveaxis(x, axis, 0), node
-    )
-    if axis is not None
-    else None,
-    options.in_kwargs_axes,
-    kwargs,
-    is_leaf=lambda x: x is None,
+      lambda axis, node: jax.tree_util.tree_map(
+          lambda x: jnp.moveaxis(x, axis, 0), node
+      )
+      if axis is not None
+      else None,
+      options.in_kwargs_axes,
+      kwargs,
+      is_leaf=lambda x: x is None,
   )
   broadcast_kwargs = jax.tree_util.tree_map(
-    lambda axis, node: None if axis is not None else node,
-    options.in_kwargs_axes,
-    kwargs,
-    is_leaf=lambda x: x is None,
+      lambda axis, node: None if axis is not None else node,
+      options.in_kwargs_axes,
+      kwargs,
+      is_leaf=lambda x: x is None,
   )
 
   # infer length
@@ -934,18 +938,18 @@ def scan_apply(
 
     # merge args and kwargs
     args = jax.tree_util.tree_map(
-      lambda axis, scan, broadcast: scan if axis is not None else broadcast,
-      options.in_args_axes,
-      scan_args,
-      broadcast_args,
-      is_leaf=lambda x: x is None,
+        lambda axis, scan, broadcast: scan if axis is not None else broadcast,
+        options.in_args_axes,
+        scan_args,
+        broadcast_args,
+        is_leaf=lambda x: x is None,
     )
     kwargs = jax.tree_util.tree_map(
-      lambda axis, scan, broadcast: scan if axis is not None else broadcast,
-      options.in_kwargs_axes,
-      scan_kwargs,
-      broadcast_kwargs,
-      is_leaf=lambda x: x is None,
+        lambda axis, scan, broadcast: scan if axis is not None else broadcast,
+        options.in_kwargs_axes,
+        scan_kwargs,
+        broadcast_kwargs,
+        is_leaf=lambda x: x is None,
     )
 
     # merge rng state
@@ -1014,16 +1018,16 @@ def scan_apply(
 
   # transpose axes state
   scan_states = tuple(
-    jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, 0, axis), axes_state)
-    for axes_state, axis in zip(scan_states, options.variable_axes.values())
+      jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, 0, axis), axes_state)
+      for axes_state, axis in zip(scan_states, options.variable_axes.values())
   )
   # transpose axes arg
   scan_out = jax.tree_util.tree_map(
-    lambda axis, node: jax.tree_util.tree_map(
-      lambda x: jnp.moveaxis(x, 0, axis), node
-    ),
-    options.out_axes,
-    scan_out,
+      lambda axis, node: jax.tree_util.tree_map(
+          lambda x: jnp.moveaxis(x, 0, axis), node
+      ),
+      options.out_axes,
+      scan_out,
   )
   # slice new carry state
   carry_state_new = jax.tree_util.tree_map(lambda x: x[0], carry_state_new)
@@ -1479,20 +1483,20 @@ def vmap_apply(
   # infer length
   axis_sizes: tp.Set[int] = set()
   args_sizes = jax.tree_util.tree_map(
-    lambda axis, node: jax.tree_util.tree_map(lambda x: x.shape[axis], node)
-    if axis is not None
-    else None,
-    options.in_args_axes,
-    args,
-    is_leaf=lambda x: x is None,
+      lambda axis, node: jax.tree_util.tree_map(lambda x: x.shape[axis], node)
+      if axis is not None
+      else None,
+      options.in_args_axes,
+      args,
+      is_leaf=lambda x: x is None,
   )
   kwargs_sizes = jax.tree_util.tree_map(
-    lambda axis, node: jax.tree_util.tree_map(lambda x: x.shape[axis], node)
-    if axis is not None
-    else None,
-    options.in_kwargs_axes,
-    kwargs,
-    is_leaf=lambda x: x is None,
+      lambda axis, node: jax.tree_util.tree_map(lambda x: x.shape[axis], node)
+      if axis is not None
+      else None,
+      options.in_kwargs_axes,
+      kwargs,
+      is_leaf=lambda x: x is None,
   )
   axis_sizes.update(jax.tree_util.tree_leaves(args_sizes))
   axis_sizes.update(jax.tree_util.tree_leaves(kwargs_sizes))
diff --git a/flax/linen/linear.py b/flax/linen/linear.py
index a7b6955c..9572b5ba 100644
--- a/flax/linen/linear.py
+++ b/flax/linen/linear.py
@@ -738,15 +738,16 @@ class ConvTranspose(Module):
       kernel size can be passed as an integer, which will be interpreted as a
       tuple of the single integer. For all other cases, it must be a sequence of
       integers.
-    strides: a sequence of `n` integers, representing the inter-window strides.
+    strides: an integer or a sequence of `n` integers, representing the
+      inter-window strides.
     padding: either the string `'SAME'`, the string `'VALID'`, the string
       `'CIRCULAR'` (periodic boundary conditions), or a sequence of `n` `(low,
       high)` integer pairs that give the padding to apply before and after each
       spatial dimension. A single int is interpreted as applying the same
       padding in all dims and assign a single int in a sequence causes the same
       padding to be used on both sides.
-    kernel_dilation: ``None``, or a sequence of ``n`` integers, giving the
-      dilation factor to apply in each spatial dimension of the convolution
+    kernel_dilation: ``None``, or an integer or a sequence of ``n`` integers,
+      giving the dilation factor to apply in each spatial dimension of the convolution
       kernel. Convolution with kernel dilation is also known as 'atrous
       convolution'.
     use_bias: whether to add a bias to the output (default: True).
@@ -804,6 +805,17 @@ class ConvTranspose(Module):
     else:
       kernel_size = tuple(self.kernel_size)
 
+    def maybe_broadcast(
+      x: Optional[Union[int, Sequence[int]]],
+    ) -> Tuple[int, ...]:
+      if x is None:
+        # backward compatibility with using None as sentinel for
+        # broadcast 1
+        x = 1
+      if isinstance(x, int):
+        return (x,) * len(kernel_size)
+      return tuple(x)
+
     # Combine all input batch dimensions into a single leading batch axis.
     num_batch_dimensions = inputs.ndim - (len(kernel_size) + 1)
     if num_batch_dimensions != 1:
@@ -814,11 +826,8 @@ class ConvTranspose(Module):
       ]
       inputs = jnp.reshape(inputs, flat_input_shape)
 
-    strides: Tuple[int, ...]
-    if self.strides is None:
-      strides = (1,) * (inputs.ndim - 2)
-    else:
-      strides = tuple(self.strides)
+    strides = maybe_broadcast(self.strides)
+    kernel_dilation = maybe_broadcast(self.kernel_dilation)
 
     in_features = jnp.shape(inputs)[-1]
     if self.transpose_kernel:
@@ -857,7 +866,7 @@ class ConvTranspose(Module):
       kernel,
       strides,
       padding_lax,
-      rhs_dilation=self.kernel_dilation,
+      rhs_dilation=kernel_dilation,
       transpose_kernel=self.transpose_kernel,
       precision=self.precision,
     )


Test patch:
diff --git a/flax/experimental/nnx/tests/test_module.py b/flax/experimental/nnx/tests/test_module.py
index c3c24de2..ad38cbbc 100644
--- a/flax/experimental/nnx/tests/test_module.py
+++ b/flax/experimental/nnx/tests/test_module.py
@@ -88,6 +88,22 @@ class TestModule:
 
     f(1.0)
 
+  def test_trace_level_error_on_nnx_grad(self):
+    # error occurs because nnx updates its nnx_trace
+    # in nnx.grad.
+    m = nnx.Dict(a=nnx.Param(1.0))
+
+    @nnx.grad
+    def f(_):
+      with pytest.raises(
+        nnx.TraceContextError,
+        match='Cannot mutate Module from different trace level',
+      ):
+        m.a = 2.0
+      return 1.0
+
+    f(m)
+
   def test_call(self):
     class Foo(nnx.Module):
       def __init__(self, c: float, *, rngs: nnx.Rngs):
diff --git a/tests/linen/linen_linear_test.py b/tests/linen/linen_linear_test.py
index b63b8fa0..3a293059 100644
--- a/tests/linen/linen_linear_test.py
+++ b/tests/linen/linen_linear_test.py
@@ -980,7 +980,7 @@ class LinearTest(parameterized.TestCase):
     conv_module = nn.ConvTranspose(
       features=4,
       use_bias=use_bias,
-      strides=(2, 2),
+      strides=2,
       kernel_size=(6, 6),
       padding='CIRCULAR',
       transpose_kernel=True,