Forward all arguments when using nnx.transforms.deprecated.scan as a …

…decorator.

PiperOrigin-RevId: 675988367

flax/nnx/transforms/deprecated.py

@@ -1238,7 +1238,19 @@ def scan(
 ) -> F | tp.Callable[[F], F]:
  if isinstance(f, Missing):
  return functools.partial(
- scan, length=length, reverse=reverse, unroll=unroll
+ scan,
+ length=length,
+ reverse=reverse,
+ unroll=unroll,
+ _split_transpose=_split_transpose,
+ in_axes=in_axes,
+ in_axes_kwargs=in_axes_kwargs,
+ out_axes=out_axes,
+ carry_argnum=carry_argnum,
+ state_axes=state_axes,
+ split_rngs=split_rngs,
+ transform_metadata=transform_metadata,
+ scan_output=scan_output,
  ) # type: ignore[return-value]
 
  @functools.wraps(f)