We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Support `unrolled` steps in `jax_utils.scan_in_dims`
Motivated by [jax#3094](https://github.com/google/jax/issues/3094), [jax#3738](https://github.com/google/jax/pull/3738) and [jax#3076](https://github.com/google/jax/pull/3076), `jax.lax.scan` currently supports specifying the number of scan iterations to unroll into a single iteration of the loop using the argument `unrolls`.

It would be nice to be able to control this from `jax_utils.scan_in_dims`.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `flax/jax_utils.py`
Content:
```
1 # Copyright 2021 The Flax Authors.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """Utilities we could consider upstreaming to Jax.
16 """
17 
18 import collections
19 from collections.abc import Iterable  # pylint: disable=g-importing-member
20 import itertools
21 import warnings
22 
23 import jax
24 from jax import lax
25 from jax import linear_util as lu
26 from jax.config import config
27 from jax.interpreters import partial_eval as pe
28 from jax.interpreters import xla
29 import jax.lib.xla_bridge as xb
30 import jax.numpy as jnp
31 import numpy as np
32 
33 
34 def replicate(tree, devices=None):
35   """Replicates arrays to multiple devices.
36 
37   Args:
38     tree: a pytree containing the arrays that should be replicated.
39     devices: the devices the data is replicated to
40       (default: `jax.local_devices()`).
41   Returns:
42     A new pytree containing the replicated arrays.
43   """
44   if devices is None:
45     # match the default device assignments used in pmap:
46     # for single-host, that's the XLA default device assignment
47     # for multi-host, it's the order of jax.local_devices()
48     if jax.process_count() == 1:
49       devices = [d for d in xb.get_backend().get_default_device_assignment(
50           jax.device_count()) if d.process_index == jax.process_index()]
51     else:
52       devices = jax.local_devices()
53 
54   return jax.device_put_replicated(tree, devices)
55 
56 
57 def unreplicate(tree):
58   """Returns a single instance of a replicated array."""
59   return jax.tree_map(lambda x: x[0], tree)
60 
61 
62 def pmean(xs, axis_name):
63   warnings.warn('use jax.lax.pmean instead',
64                 DeprecationWarning)
65   return lax.pmean(xs, axis_name)
66 
67 
68 def partial_eval_by_shape(fn, input_spec, *args, **kwargs):
69   """Lazily evaluate a function by using the shapes of the inputs.
70 
71   This function is similar to `jax.eval_shape` with the key difference that
72   function outputs that can be computed without a concrete value of the
73   inputs are returned as is instead of only the shape. See for example
74   `module.init_by_shape` where this functionality is used to initialize a
75   model without using input data lr computation.
76 
77   Args:
78     fn: the function to be lazily evaluated.
79     input_spec: an iterable of shapes or (shape, dtype) tuples specifying the
80       shape and type of the inputs. If unspecified the dtype is float32.
81     *args: other arguments passed to the module's apply function
82     **kwargs: keyword arguments passed to the module's apply function
83   Returns:
84     A pair consisting of the model output and an instance of Model
85   """
86   # output cannot be returned in lazy_create because jax.eval_shape will only
87   # return the shape and dtype.
88   # TODO(mattjj,jheek): use a public JAX API
89   f = lambda *inputs: fn(*inputs, *args, **kwargs)
90   input_structs = [_parse_spec(spec) for spec in input_spec]
91   inputs_flat, in_tree = jax.tree_flatten(input_structs)
92   f_flat, out_tree = jax.api_util.flatten_fun_nokwargs(lu.wrap_init(f), in_tree)
93   in_pvals = [pe.PartialVal.unknown(jax.ShapedArray(x.shape, x.dtype))
94               for x in inputs_flat]
95   _, out_pvals, _ = pe.trace_to_jaxpr(f_flat, in_pvals)
96   out_flat = [const if pv is None else jax.ShapeDtypeStruct(pv.shape, pv.dtype)
97               for pv, const in out_pvals]
98   return jax.tree_unflatten(out_tree(), out_flat)
99 
100 
101 def _parse_spec(spec):
102   """Parse an input spec of the form (shape, dtype) or shape into a jax.ShapeDtypeStruct."""
103   spec = tuple(spec)
104   if len(spec) == 2 and isinstance(spec[0], Iterable):
105     return jax.ShapeDtypeStruct(tuple(spec[0]), spec[1])
106   else:
107     return jax.ShapeDtypeStruct(spec, jnp.float32)
108 
109 
110 def prefetch_to_device(iterator, size, devices=None):
111   """Shard and prefetch batches on device.
112 
113   This utility takes an iterator and returns a new iterator which fills an on
114   device prefetch buffer. Eager prefetching can improve the performance of
115   training loops significantly by overlapping compute and data transfer.
116 
117   This utility is mostly useful for GPUs, for TPUs and CPUs it should not be
118   necessary -- the TPU & CPU memory allocators (normally) don't pick a memory
119   location that isn't free yet so they don't block. Instead those allocators OOM.
120 
121   Args:
122     iterator: an iterator that yields a pytree of ndarrays where the first
123       dimension is sharded across devices.
124 
125     size: the size of the prefetch buffer.
126 
127       If you're training on GPUs, 2 is generally the best choice because this
128       guarantees that you can overlap a training step on GPU with a data
129       prefetch step on CPU.
130 
131     devices: the list of devices to which the arrays should be prefetched.
132 
133   Yields:
134     The original items from the iterator where each ndarray is now a sharded to
135     the specified devices.
136   """
137   queue = collections.deque()
138   devices = devices or jax.local_devices()
139 
140   def _prefetch(xs):
141     if hasattr(jax, "device_put_sharded"):  # jax>=0.2.0
142       return jax.device_put_sharded(list(xs), devices)
143     else:
144       aval = jax.xla.abstractify(xs)
145       assert xs.shape[0] == len(devices), (
146           "The first dimension of the iterator's ndarrays is not "
147           "equal to the number of devices.")
148       buffers = [xla.device_put(x, devices[i])
149                  for i, x in enumerate(xs)]
150       return jax.pxla.ShardedDeviceArray(aval, buffers)
151 
152   def enqueue(n):  # Enqueues *up to* `n` elements from the iterator.
153     for data in itertools.islice(iterator, n):
154       queue.append(jax.tree_map(_prefetch, data))
155 
156   enqueue(size)  # Fill up the buffer.
157   while queue:
158     yield queue.popleft()
159     enqueue(1)
160 
161 
162 def _scan_nd(body_fn, init, xs, n=1):
163   """Utility for performing an n-dimensional `lax.scan`.
164 
165   The n-d scan is simply recursive call of 1-d scan.
166   Args:
167     body_fn: the body of the loop of type (c, x) -> (c, y).
168     init: initial value for the carry.
169     xs: a pytree of tensors to scan over.
170     n: number of dimensions to scan over (default: 1)
171   Returns:
172     A tuple of the final carry and the values returned by the body.
173   """
174   if n == 1:
175     return lax.scan(body_fn, init, xs)
176   else:
177     def scan_body(c, x):
178       return _scan_nd(body_fn, c, x, n=n-1)
179     return lax.scan(scan_body, init, xs)
180 
181 
182 def _invert_perm(perm):
183   perm_inv = [0] * len(perm)
184   for i, j in enumerate(perm):
185     perm_inv[j] = i
186   return tuple(perm_inv)
187 
188 
189 def scan_in_dim(body_fn, init, xs, axis=(0,), keepdims=False):
190   """utility for doing a scan along arbitrary dimensions.
191 
192   see `lax.scan` for details on how the scan operation works.
193   Args:
194     body_fn: the body of the loop of type (c, x) -> (c, y).
195     init: initial value for the carry.
196     xs: a pytree of tensors to scan over.
197     axis: the axis to scan over.
198     keepdims: keep the dimensions that are scanned over.
199   Returns:
200     A tuple of the final carry and the values returned by the body.
201   """
202   if not isinstance(axis, Iterable):
203     axis = (axis,)
204 
205   def transpose_in(x):
206     perm = axis + tuple(np.delete(np.arange(x.ndim), axis))
207     return x.transpose(perm)
208   def transpose_out(x):
209     perm = axis + tuple(np.delete(np.arange(x.ndim), axis))
210     return x.transpose(_invert_perm(perm))
211 
212   def body_wrapper(c, xs):
213     if keepdims:
214       xs = jax.tree_map(lambda x: x.reshape((1,) * len(axis) + x.shape), xs)
215       xs = jax.tree_map(transpose_out, xs)
216     c, ys = body_fn(c, xs)
217     if keepdims:
218       ys = jax.tree_map(transpose_in, ys)
219       ys = jax.tree_map(lambda x: x.reshape(x.shape[len(axis):]), ys)
220     return c, ys
221 
222   xs = jax.tree_map(transpose_in, xs)
223   c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis))
224   ys = jax.tree_map(transpose_out, ys)
225   return c, ys
226
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/flax/jax_utils.py b/flax/jax_utils.py
--- a/flax/jax_utils.py
+++ b/flax/jax_utils.py
@@ -159,7 +159,7 @@
     enqueue(1)
 
 
-def _scan_nd(body_fn, init, xs, n=1):
+def _scan_nd(body_fn, init, xs, n=1, unroll=(1,)):
   """Utility for performing an n-dimensional `lax.scan`.
 
   The n-d scan is simply recursive call of 1-d scan.
@@ -172,11 +172,11 @@
     A tuple of the final carry and the values returned by the body.
   """
   if n == 1:
-    return lax.scan(body_fn, init, xs)
+    return lax.scan(body_fn, init, xs, unroll=unroll[0])
   else:
     def scan_body(c, x):
-      return _scan_nd(body_fn, c, x, n=n-1)
-    return lax.scan(scan_body, init, xs)
+      return _scan_nd(body_fn, c, x, n=n-1, unroll=unroll[1:])
+    return lax.scan(scan_body, init, xs, unroll=unroll[0])
 
 
 def _invert_perm(perm):
@@ -186,22 +186,38 @@
   return tuple(perm_inv)
 
 
-def scan_in_dim(body_fn, init, xs, axis=(0,), keepdims=False):
+def scan_in_dim(body_fn, init, xs, axis=(0,), unroll=(1,), keepdims=False):
   """utility for doing a scan along arbitrary dimensions.
 
-  see `lax.scan` for details on how the scan operation works.
+  See `lax.scan` for details on how the scan operation works.
+
+  Note on `unroll`: This argument gets left padded with ones to match the size
+  of `axis`. Doing so allows unrolls to performed from the innermost loop first.
+  For example, `scan_in_dim(..., axis=(1, 2, 3), unroll=5)` is equivalent to
+  `scan_in_dim(..., axis=(1, 2, 3), unroll=(1, 1, 5))`.
+
   Args:
     body_fn: the body of the loop of type (c, x) -> (c, y).
     init: initial value for the carry.
     xs: a pytree of tensors to scan over.
     axis: the axis to scan over.
     keepdims: keep the dimensions that are scanned over.
+    unroll: an optional positive integer, or tuple of positive integers
+      showing how many iterations of the loop to be unroll into a single
+      iteration for each axis.
   Returns:
     A tuple of the final carry and the values returned by the body.
   """
   if not isinstance(axis, Iterable):
     axis = (axis,)
 
+  if not isinstance(unroll, Iterable):
+    unroll = (unroll,)
+
+  # Pad unroll with ones so we start unrolling from the innermost loop
+  len_diff = len(axis) - len(unroll)
+  unroll = (1,) * len_diff + unroll
+
   def transpose_in(x):
     perm = axis + tuple(np.delete(np.arange(x.ndim), axis))
     return x.transpose(perm)
@@ -220,6 +236,6 @@
     return c, ys
 
   xs = jax.tree_map(transpose_in, xs)
-  c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis))
+  c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis), unroll=unroll)
   ys = jax.tree_map(transpose_out, ys)
   return c, ys

{"golden_diff": "diff --git a/flax/jax_utils.py b/flax/jax_utils.py\n--- a/flax/jax_utils.py\n+++ b/flax/jax_utils.py\n@@ -159,7 +159,7 @@\n     enqueue(1)\n \n \n-def _scan_nd(body_fn, init, xs, n=1):\n+def _scan_nd(body_fn, init, xs, n=1, unroll=(1,)):\n   \"\"\"Utility for performing an n-dimensional `lax.scan`.\n \n   The n-d scan is simply recursive call of 1-d scan.\n@@ -172,11 +172,11 @@\n     A tuple of the final carry and the values returned by the body.\n   \"\"\"\n   if n == 1:\n-    return lax.scan(body_fn, init, xs)\n+    return lax.scan(body_fn, init, xs, unroll=unroll[0])\n   else:\n     def scan_body(c, x):\n-      return _scan_nd(body_fn, c, x, n=n-1)\n-    return lax.scan(scan_body, init, xs)\n+      return _scan_nd(body_fn, c, x, n=n-1, unroll=unroll[1:])\n+    return lax.scan(scan_body, init, xs, unroll=unroll[0])\n \n \n def _invert_perm(perm):\n@@ -186,22 +186,38 @@\n   return tuple(perm_inv)\n \n \n-def scan_in_dim(body_fn, init, xs, axis=(0,), keepdims=False):\n+def scan_in_dim(body_fn, init, xs, axis=(0,), unroll=(1,), keepdims=False):\n   \"\"\"utility for doing a scan along arbitrary dimensions.\n \n-  see `lax.scan` for details on how the scan operation works.\n+  See `lax.scan` for details on how the scan operation works.\n+\n+  Note on `unroll`: This argument gets left padded with ones to match the size\n+  of `axis`. Doing so allows unrolls to performed from the innermost loop first.\n+  For example, `scan_in_dim(..., axis=(1, 2, 3), unroll=5)` is equivalent to\n+  `scan_in_dim(..., axis=(1, 2, 3), unroll=(1, 1, 5))`.\n+\n   Args:\n     body_fn: the body of the loop of type (c, x) -> (c, y).\n     init: initial value for the carry.\n     xs: a pytree of tensors to scan over.\n     axis: the axis to scan over.\n     keepdims: keep the dimensions that are scanned over.\n+    unroll: an optional positive integer, or tuple of positive integers\n+      showing how many iterations of the loop to be unroll into a single\n+      iteration for each axis.\n   Returns:\n     A tuple of the final carry and the values returned by the body.\n   \"\"\"\n   if not isinstance(axis, Iterable):\n     axis = (axis,)\n \n+  if not isinstance(unroll, Iterable):\n+    unroll = (unroll,)\n+\n+  # Pad unroll with ones so we start unrolling from the innermost loop\n+  len_diff = len(axis) - len(unroll)\n+  unroll = (1,) * len_diff + unroll\n+\n   def transpose_in(x):\n     perm = axis + tuple(np.delete(np.arange(x.ndim), axis))\n     return x.transpose(perm)\n@@ -220,6 +236,6 @@\n     return c, ys\n \n   xs = jax.tree_map(transpose_in, xs)\n-  c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis))\n+  c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis), unroll=unroll)\n   ys = jax.tree_map(transpose_out, ys)\n   return c, ys\n", "issue": "Support `unrolled` steps in `jax_utils.scan_in_dims`\nMotivated by [jax#3094](https://github.com/google/jax/issues/3094), [jax#3738](https://github.com/google/jax/pull/3738) and [jax#3076](https://github.com/google/jax/pull/3076), `jax.lax.scan` currently supports specifying the number of scan iterations to unroll into a single iteration of the loop using the argument `unrolls`.\r\n\r\nIt would be nice to be able to control this from `jax_utils.scan_in_dims`.\n", "before_files": [{"content": "# Copyright 2021 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities we could consider upstreaming to Jax.\n\"\"\"\n\nimport collections\nfrom collections.abc import Iterable  # pylint: disable=g-importing-member\nimport itertools\nimport warnings\n\nimport jax\nfrom jax import lax\nfrom jax import linear_util as lu\nfrom jax.config import config\nfrom jax.interpreters import partial_eval as pe\nfrom jax.interpreters import xla\nimport jax.lib.xla_bridge as xb\nimport jax.numpy as jnp\nimport numpy as np\n\n\ndef replicate(tree, devices=None):\n  \"\"\"Replicates arrays to multiple devices.\n\n  Args:\n    tree: a pytree containing the arrays that should be replicated.\n    devices: the devices the data is replicated to\n      (default: `jax.local_devices()`).\n  Returns:\n    A new pytree containing the replicated arrays.\n  \"\"\"\n  if devices is None:\n    # match the default device assignments used in pmap:\n    # for single-host, that's the XLA default device assignment\n    # for multi-host, it's the order of jax.local_devices()\n    if jax.process_count() == 1:\n      devices = [d for d in xb.get_backend().get_default_device_assignment(\n          jax.device_count()) if d.process_index == jax.process_index()]\n    else:\n      devices = jax.local_devices()\n\n  return jax.device_put_replicated(tree, devices)\n\n\ndef unreplicate(tree):\n  \"\"\"Returns a single instance of a replicated array.\"\"\"\n  return jax.tree_map(lambda x: x[0], tree)\n\n\ndef pmean(xs, axis_name):\n  warnings.warn('use jax.lax.pmean instead',\n                DeprecationWarning)\n  return lax.pmean(xs, axis_name)\n\n\ndef partial_eval_by_shape(fn, input_spec, *args, **kwargs):\n  \"\"\"Lazily evaluate a function by using the shapes of the inputs.\n\n  This function is similar to `jax.eval_shape` with the key difference that\n  function outputs that can be computed without a concrete value of the\n  inputs are returned as is instead of only the shape. See for example\n  `module.init_by_shape` where this functionality is used to initialize a\n  model without using input data lr computation.\n\n  Args:\n    fn: the function to be lazily evaluated.\n    input_spec: an iterable of shapes or (shape, dtype) tuples specifying the\n      shape and type of the inputs. If unspecified the dtype is float32.\n    *args: other arguments passed to the module's apply function\n    **kwargs: keyword arguments passed to the module's apply function\n  Returns:\n    A pair consisting of the model output and an instance of Model\n  \"\"\"\n  # output cannot be returned in lazy_create because jax.eval_shape will only\n  # return the shape and dtype.\n  # TODO(mattjj,jheek): use a public JAX API\n  f = lambda *inputs: fn(*inputs, *args, **kwargs)\n  input_structs = [_parse_spec(spec) for spec in input_spec]\n  inputs_flat, in_tree = jax.tree_flatten(input_structs)\n  f_flat, out_tree = jax.api_util.flatten_fun_nokwargs(lu.wrap_init(f), in_tree)\n  in_pvals = [pe.PartialVal.unknown(jax.ShapedArray(x.shape, x.dtype))\n              for x in inputs_flat]\n  _, out_pvals, _ = pe.trace_to_jaxpr(f_flat, in_pvals)\n  out_flat = [const if pv is None else jax.ShapeDtypeStruct(pv.shape, pv.dtype)\n              for pv, const in out_pvals]\n  return jax.tree_unflatten(out_tree(), out_flat)\n\n\ndef _parse_spec(spec):\n  \"\"\"Parse an input spec of the form (shape, dtype) or shape into a jax.ShapeDtypeStruct.\"\"\"\n  spec = tuple(spec)\n  if len(spec) == 2 and isinstance(spec[0], Iterable):\n    return jax.ShapeDtypeStruct(tuple(spec[0]), spec[1])\n  else:\n    return jax.ShapeDtypeStruct(spec, jnp.float32)\n\n\ndef prefetch_to_device(iterator, size, devices=None):\n  \"\"\"Shard and prefetch batches on device.\n\n  This utility takes an iterator and returns a new iterator which fills an on\n  device prefetch buffer. Eager prefetching can improve the performance of\n  training loops significantly by overlapping compute and data transfer.\n\n  This utility is mostly useful for GPUs, for TPUs and CPUs it should not be\n  necessary -- the TPU & CPU memory allocators (normally) don't pick a memory\n  location that isn't free yet so they don't block. Instead those allocators OOM.\n\n  Args:\n    iterator: an iterator that yields a pytree of ndarrays where the first\n      dimension is sharded across devices.\n\n    size: the size of the prefetch buffer.\n\n      If you're training on GPUs, 2 is generally the best choice because this\n      guarantees that you can overlap a training step on GPU with a data\n      prefetch step on CPU.\n\n    devices: the list of devices to which the arrays should be prefetched.\n\n  Yields:\n    The original items from the iterator where each ndarray is now a sharded to\n    the specified devices.\n  \"\"\"\n  queue = collections.deque()\n  devices = devices or jax.local_devices()\n\n  def _prefetch(xs):\n    if hasattr(jax, \"device_put_sharded\"):  # jax>=0.2.0\n      return jax.device_put_sharded(list(xs), devices)\n    else:\n      aval = jax.xla.abstractify(xs)\n      assert xs.shape[0] == len(devices), (\n          \"The first dimension of the iterator's ndarrays is not \"\n          \"equal to the number of devices.\")\n      buffers = [xla.device_put(x, devices[i])\n                 for i, x in enumerate(xs)]\n      return jax.pxla.ShardedDeviceArray(aval, buffers)\n\n  def enqueue(n):  # Enqueues *up to* `n` elements from the iterator.\n    for data in itertools.islice(iterator, n):\n      queue.append(jax.tree_map(_prefetch, data))\n\n  enqueue(size)  # Fill up the buffer.\n  while queue:\n    yield queue.popleft()\n    enqueue(1)\n\n\ndef _scan_nd(body_fn, init, xs, n=1):\n  \"\"\"Utility for performing an n-dimensional `lax.scan`.\n\n  The n-d scan is simply recursive call of 1-d scan.\n  Args:\n    body_fn: the body of the loop of type (c, x) -> (c, y).\n    init: initial value for the carry.\n    xs: a pytree of tensors to scan over.\n    n: number of dimensions to scan over (default: 1)\n  Returns:\n    A tuple of the final carry and the values returned by the body.\n  \"\"\"\n  if n == 1:\n    return lax.scan(body_fn, init, xs)\n  else:\n    def scan_body(c, x):\n      return _scan_nd(body_fn, c, x, n=n-1)\n    return lax.scan(scan_body, init, xs)\n\n\ndef _invert_perm(perm):\n  perm_inv = [0] * len(perm)\n  for i, j in enumerate(perm):\n    perm_inv[j] = i\n  return tuple(perm_inv)\n\n\ndef scan_in_dim(body_fn, init, xs, axis=(0,), keepdims=False):\n  \"\"\"utility for doing a scan along arbitrary dimensions.\n\n  see `lax.scan` for details on how the scan operation works.\n  Args:\n    body_fn: the body of the loop of type (c, x) -> (c, y).\n    init: initial value for the carry.\n    xs: a pytree of tensors to scan over.\n    axis: the axis to scan over.\n    keepdims: keep the dimensions that are scanned over.\n  Returns:\n    A tuple of the final carry and the values returned by the body.\n  \"\"\"\n  if not isinstance(axis, Iterable):\n    axis = (axis,)\n\n  def transpose_in(x):\n    perm = axis + tuple(np.delete(np.arange(x.ndim), axis))\n    return x.transpose(perm)\n  def transpose_out(x):\n    perm = axis + tuple(np.delete(np.arange(x.ndim), axis))\n    return x.transpose(_invert_perm(perm))\n\n  def body_wrapper(c, xs):\n    if keepdims:\n      xs = jax.tree_map(lambda x: x.reshape((1,) * len(axis) + x.shape), xs)\n      xs = jax.tree_map(transpose_out, xs)\n    c, ys = body_fn(c, xs)\n    if keepdims:\n      ys = jax.tree_map(transpose_in, ys)\n      ys = jax.tree_map(lambda x: x.reshape(x.shape[len(axis):]), ys)\n    return c, ys\n\n  xs = jax.tree_map(transpose_in, xs)\n  c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis))\n  ys = jax.tree_map(transpose_out, ys)\n  return c, ys\n", "path": "flax/jax_utils.py"}], "after_files": [{"content": "# Copyright 2021 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities we could consider upstreaming to Jax.\n\"\"\"\n\nimport collections\nfrom collections.abc import Iterable  # pylint: disable=g-importing-member\nimport itertools\nimport warnings\n\nimport jax\nfrom jax import lax\nfrom jax import linear_util as lu\nfrom jax.config import config\nfrom jax.interpreters import partial_eval as pe\nfrom jax.interpreters import xla\nimport jax.lib.xla_bridge as xb\nimport jax.numpy as jnp\nimport numpy as np\n\n\ndef replicate(tree, devices=None):\n  \"\"\"Replicates arrays to multiple devices.\n\n  Args:\n    tree: a pytree containing the arrays that should be replicated.\n    devices: the devices the data is replicated to\n      (default: `jax.local_devices()`).\n  Returns:\n    A new pytree containing the replicated arrays.\n  \"\"\"\n  if devices is None:\n    # match the default device assignments used in pmap:\n    # for single-host, that's the XLA default device assignment\n    # for multi-host, it's the order of jax.local_devices()\n    if jax.process_count() == 1:\n      devices = [d for d in xb.get_backend().get_default_device_assignment(\n          jax.device_count()) if d.process_index == jax.process_index()]\n    else:\n      devices = jax.local_devices()\n\n  return jax.device_put_replicated(tree, devices)\n\n\ndef unreplicate(tree):\n  \"\"\"Returns a single instance of a replicated array.\"\"\"\n  return jax.tree_map(lambda x: x[0], tree)\n\n\ndef pmean(xs, axis_name):\n  warnings.warn('use jax.lax.pmean instead',\n                DeprecationWarning)\n  return lax.pmean(xs, axis_name)\n\n\ndef partial_eval_by_shape(fn, input_spec, *args, **kwargs):\n  \"\"\"Lazily evaluate a function by using the shapes of the inputs.\n\n  This function is similar to `jax.eval_shape` with the key difference that\n  function outputs that can be computed without a concrete value of the\n  inputs are returned as is instead of only the shape. See for example\n  `module.init_by_shape` where this functionality is used to initialize a\n  model without using input data lr computation.\n\n  Args:\n    fn: the function to be lazily evaluated.\n    input_spec: an iterable of shapes or (shape, dtype) tuples specifying the\n      shape and type of the inputs. If unspecified the dtype is float32.\n    *args: other arguments passed to the module's apply function\n    **kwargs: keyword arguments passed to the module's apply function\n  Returns:\n    A pair consisting of the model output and an instance of Model\n  \"\"\"\n  # output cannot be returned in lazy_create because jax.eval_shape will only\n  # return the shape and dtype.\n  # TODO(mattjj,jheek): use a public JAX API\n  f = lambda *inputs: fn(*inputs, *args, **kwargs)\n  input_structs = [_parse_spec(spec) for spec in input_spec]\n  inputs_flat, in_tree = jax.tree_flatten(input_structs)\n  f_flat, out_tree = jax.api_util.flatten_fun_nokwargs(lu.wrap_init(f), in_tree)\n  in_pvals = [pe.PartialVal.unknown(jax.ShapedArray(x.shape, x.dtype))\n              for x in inputs_flat]\n  _, out_pvals, _ = pe.trace_to_jaxpr(f_flat, in_pvals)\n  out_flat = [const if pv is None else jax.ShapeDtypeStruct(pv.shape, pv.dtype)\n              for pv, const in out_pvals]\n  return jax.tree_unflatten(out_tree(), out_flat)\n\n\ndef _parse_spec(spec):\n  \"\"\"Parse an input spec of the form (shape, dtype) or shape into a jax.ShapeDtypeStruct.\"\"\"\n  spec = tuple(spec)\n  if len(spec) == 2 and isinstance(spec[0], Iterable):\n    return jax.ShapeDtypeStruct(tuple(spec[0]), spec[1])\n  else:\n    return jax.ShapeDtypeStruct(spec, jnp.float32)\n\n\ndef prefetch_to_device(iterator, size, devices=None):\n  \"\"\"Shard and prefetch batches on device.\n\n  This utility takes an iterator and returns a new iterator which fills an on\n  device prefetch buffer. Eager prefetching can improve the performance of\n  training loops significantly by overlapping compute and data transfer.\n\n  This utility is mostly useful for GPUs, for TPUs and CPUs it should not be\n  necessary -- the TPU & CPU memory allocators (normally) don't pick a memory\n  location that isn't free yet so they don't block. Instead those allocators OOM.\n\n  Args:\n    iterator: an iterator that yields a pytree of ndarrays where the first\n      dimension is sharded across devices.\n\n    size: the size of the prefetch buffer.\n\n      If you're training on GPUs, 2 is generally the best choice because this\n      guarantees that you can overlap a training step on GPU with a data\n      prefetch step on CPU.\n\n    devices: the list of devices to which the arrays should be prefetched.\n\n  Yields:\n    The original items from the iterator where each ndarray is now a sharded to\n    the specified devices.\n  \"\"\"\n  queue = collections.deque()\n  devices = devices or jax.local_devices()\n\n  def _prefetch(xs):\n    if hasattr(jax, \"device_put_sharded\"):  # jax>=0.2.0\n      return jax.device_put_sharded(list(xs), devices)\n    else:\n      aval = jax.xla.abstractify(xs)\n      assert xs.shape[0] == len(devices), (\n          \"The first dimension of the iterator's ndarrays is not \"\n          \"equal to the number of devices.\")\n      buffers = [xla.device_put(x, devices[i])\n                 for i, x in enumerate(xs)]\n      return jax.pxla.ShardedDeviceArray(aval, buffers)\n\n  def enqueue(n):  # Enqueues *up to* `n` elements from the iterator.\n    for data in itertools.islice(iterator, n):\n      queue.append(jax.tree_map(_prefetch, data))\n\n  enqueue(size)  # Fill up the buffer.\n  while queue:\n    yield queue.popleft()\n    enqueue(1)\n\n\ndef _scan_nd(body_fn, init, xs, n=1, unroll=(1,)):\n  \"\"\"Utility for performing an n-dimensional `lax.scan`.\n\n  The n-d scan is simply recursive call of 1-d scan.\n  Args:\n    body_fn: the body of the loop of type (c, x) -> (c, y).\n    init: initial value for the carry.\n    xs: a pytree of tensors to scan over.\n    n: number of dimensions to scan over (default: 1)\n  Returns:\n    A tuple of the final carry and the values returned by the body.\n  \"\"\"\n  if n == 1:\n    return lax.scan(body_fn, init, xs, unroll=unroll[0])\n  else:\n    def scan_body(c, x):\n      return _scan_nd(body_fn, c, x, n=n-1, unroll=unroll[1:])\n    return lax.scan(scan_body, init, xs, unroll=unroll[0])\n\n\ndef _invert_perm(perm):\n  perm_inv = [0] * len(perm)\n  for i, j in enumerate(perm):\n    perm_inv[j] = i\n  return tuple(perm_inv)\n\n\ndef scan_in_dim(body_fn, init, xs, axis=(0,), unroll=(1,), keepdims=False):\n  \"\"\"utility for doing a scan along arbitrary dimensions.\n\n  See `lax.scan` for details on how the scan operation works.\n\n  Note on `unroll`: This argument gets left padded with ones to match the size\n  of `axis`. Doing so allows unrolls to performed from the innermost loop first.\n  For example, `scan_in_dim(..., axis=(1, 2, 3), unroll=5)` is equivalent to\n  `scan_in_dim(..., axis=(1, 2, 3), unroll=(1, 1, 5))`.\n\n  Args:\n    body_fn: the body of the loop of type (c, x) -> (c, y).\n    init: initial value for the carry.\n    xs: a pytree of tensors to scan over.\n    axis: the axis to scan over.\n    keepdims: keep the dimensions that are scanned over.\n    unroll: an optional positive integer, or tuple of positive integers\n      showing how many iterations of the loop to be unroll into a single\n      iteration for each axis.\n  Returns:\n    A tuple of the final carry and the values returned by the body.\n  \"\"\"\n  if not isinstance(axis, Iterable):\n    axis = (axis,)\n\n  if not isinstance(unroll, Iterable):\n    unroll = (unroll,)\n\n  # Pad unroll with ones so we start unrolling from the innermost loop\n  len_diff = len(axis) - len(unroll)\n  unroll = (1,) * len_diff + unroll\n\n  def transpose_in(x):\n    perm = axis + tuple(np.delete(np.arange(x.ndim), axis))\n    return x.transpose(perm)\n  def transpose_out(x):\n    perm = axis + tuple(np.delete(np.arange(x.ndim), axis))\n    return x.transpose(_invert_perm(perm))\n\n  def body_wrapper(c, xs):\n    if keepdims:\n      xs = jax.tree_map(lambda x: x.reshape((1,) * len(axis) + x.shape), xs)\n      xs = jax.tree_map(transpose_out, xs)\n    c, ys = body_fn(c, xs)\n    if keepdims:\n      ys = jax.tree_map(transpose_in, ys)\n      ys = jax.tree_map(lambda x: x.reshape(x.shape[len(axis):]), ys)\n    return c, ys\n\n  xs = jax.tree_map(transpose_in, xs)\n  c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis), unroll=unroll)\n  ys = jax.tree_map(transpose_out, ys)\n  return c, ys\n", "path": "flax/jax_utils.py"}]}