Add 'default_clobber' mode that lets different loop iterations sow di…

…fferent values to be reaped.

PiperOrigin-RevId: 612253581

oryx/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/bijectors/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/harvest_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -728,6 +728,33 @@ def f(init):
     self.assertListEqual(['x'], list(variables.keys()))
     np.testing.assert_allclose(variables['x'], true_out[-1])
 
+  @parameterized.named_parameters(('scan', True), ('while_loop', False))
+  def test_can_reap_and_plant_looped_values_in_default_clobber_mode(
+      self, static_length
+  ):
+    length = 5
+
+    def body(index, x):
+      x = x + index
+      x = jax.lax.switch(
+          index,
+          [
+              lambda i=i: variable(x, name=f'x{i}', mode='default_clobber')
+              for i in range(length + 1)
+          ],
+      )
+      return x
+
+    def f(upper, init):
+      return lax.fori_loop(0, length if static_length else upper, body, init)
+
+    out, variables = harvest_variables(f)(dict(x3=0.5), length, 1.)
+    np.testing.assert_allclose(out, 0.5 + 4)
+    default = 0.
+    self.assertDictEqual(
+        dict(x0=1., x1=2., x2=4., x4=out, x5=default), variables
+    )
+
   def test_non_clobber_mode_in_while_loop_should_error_with_reap_and_plant(
       self):
 
@@ -844,11 +871,11 @@ def false_fun(x):
       return lax.cond(pred, true_fun, false_fun, x)
 
     with self.assertRaisesRegex(
-        ValueError, 'Mismatching number of `sow`s between branches.'):
+        ValueError, 'Missing sow in branch: \'y\''):
       reap_variables(f2)(True, 1.)
 
     with self.assertRaisesRegex(
-        ValueError, 'Mismatching number of `sow`s between branches.'):
+        ValueError, 'Missing sow in branch: \'y\''):
       plant_variables(f2)({}, True, 1.)
 
     def f3(pred, x):
@@ -893,6 +920,28 @@ def false_fun(x):
     self.assertEqual(out, 6.)
     self.assertDictEqual(reaps, dict(x=3.))
 
+  def test_can_reap_from_mismatching_branches_of_default_clobber_cond(self):
+
+    def f(pred, x):
+
+      def true_fun(x):
+        x = variable(x, name='x', mode='default_clobber')
+        return x + 2.
+
+      def false_fun(x):
+        x = variable(x + 2., name='y', mode='default_clobber')
+        return x + 3.
+
+      return lax.cond(pred, true_fun, false_fun, x)
+
+    out, reaps = call_and_reap_variables(f)(True, 1.)
+    self.assertEqual(out, 3.)
+    self.assertDictEqual(reaps, dict(x=1., y=0.))
+
+    out, reaps = call_and_reap_variables(f)(False, 1.)
+    self.assertEqual(out, 6.)
+    self.assertDictEqual(reaps, dict(y=3., x=0.))
+
   def test_can_plant_values_into_either_branch_of_cond(self):
 
     def f(pred, x):
@@ -913,6 +962,26 @@ def false_fun(x):
     out = plant_variables(f)(dict(x=4.), False, 1.)
     self.assertEqual(out, 7.)
 
+  def test_can_plant_into_mismatching_branches_of_default_clobber_cond(self):
+
+    def f(pred, x):
+
+      def true_fun(x):
+        x = variable(x, name='x', mode='default_clobber')
+        return x + 2.
+
+      def false_fun(x):
+        x = variable(x + 2., name='y', mode='default_clobber')
+        return x + 3.
+
+      return lax.cond(pred, true_fun, false_fun, x)
+
+    out = plant_variables(f)(dict(x=4.), True, 1.)
+    self.assertEqual(out, 6.)
+
+    out = plant_variables(f)(dict(y=4.), False, 1.)
+    self.assertEqual(out, 7.)
+
   def test_can_reap_values_from_any_branch_in_switch(self):
 
     def f(index, x):

oryx/core/interpreters/inverse/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/inverse/bijector_extensions.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/inverse/bijector_extensions_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/inverse/core.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/inverse/custom_inverse.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/inverse/custom_inverse_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/inverse/inverse_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/inverse/rules.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/inverse/slice.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/inverse/slice_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/log_prob.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/log_prob_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/propagate.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/interpreters/propagate_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/kwargs_util.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/kwargs_util_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/ppl/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/ppl/effect_handler.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/ppl/effect_handler_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/ppl/plate_util.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/ppl/transformations.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/ppl/transformations_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/primitive.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/primitive_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/pytree.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/serialize.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/serialize_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/state/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/state/api.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/state/function.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/state/function_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/state/module.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/state/registrations.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/state/registrations_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/core/trace_util.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/distributions/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/distributions/distribution_extensions.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/distributions/distribution_extensions_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/experimental/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/experimental/autoconj/addn.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/experimental/autoconj/addn_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/experimental/autoconj/canonicalize.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/experimental/autoconj/canonicalize_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

oryx/experimental/autoconj/einsum.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The oryx Authors.
+# Copyright 2024 The oryx Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.