Add non-jittable fori_loop

jaxopt/_src/block_cd.py

@@ -28,6 +28,7 @@
 import jax.numpy as jnp
 
 from jaxopt._src import base
+from jaxopt._src.fori_loop import fori_loop
 from jaxopt._src import implicit_diff as idf
 from jaxopt._src import loop
 from jaxopt._src import objective
@@ -155,8 +156,8 @@ def body_fun(i, tup):
     # a for loop that can be potentially non-jitted.
     # this will allow to unit test the number of function eval.
     # (zramzi)
-    params, subfun_g, predictions, sqerror_sum = jax.lax.fori_loop(
-        lower=0, upper=n_for, body_fun=body_fun, init_val=init)
+    params, subfun_g, predictions, sqerror_sum = fori_loop(
+        lower=0, upper=n_for, body_fun=body_fun, init_val=init, jit=self.jit)
     state = BlockCDState(iter_num=state.iter_num + 1,
                          predictions=predictions,
                          subfun_g=subfun_g,

jaxopt/_src/fori_loop.py

@@ -0,0 +1,27 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Looping utilities."""
+
+import jax
+
+
+def fori_loop(lower, upper, body_fun, init_val, jit=True):
+    """Wrapper to avoid having the condition to be compiled if not wanted."""
+    if not jit:
+        with jax.disable_jit():
+            return jax.lax.fori_loop(
+                lower=lower, upper=upper, body_fun=body_fun, init_val=init_val)
+    return jax.lax.fori_loop(
+        lower=lower, upper=upper, body_fun=body_fun, init_val=init_val)

jaxopt/fori_loop.py

@@ -0,0 +1,15 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from jaxopt._src.fori_loop import fori_loop