Repository: keras-team/keras (python)

Issue:
Error in masked BatchNormalization with > 3 dimensions
```
import keras

x = keras.ops.ones((1, 2, 3, 4))
mask = keras.ops.ones((1, 2), dtype="bool")
y = keras.layers.BatchNormalization()(x, mask=mask, training=True)
```
gives
```
tensorflow.python.framework.errors_impl.InvalidArgumentError: Exception encountered when calling BatchNormalization.call().

{{function_node __wrapped__Mul_device_/job:localhost/replica:0/task:0/device:CPU:0}} Incompatible shapes: [1,2,1] vs. [1,2,3,4] [Op:Mul] name: 

Arguments received by BatchNormalization.call():
  • inputs=tf.Tensor(shape=(1, 2, 3, 4), dtype=float32)
  • training=True
  • mask=tf.Tensor(shape=(1, 2), dtype=bool)
```

Pull request:
Make batch norm mask shape error more descriptive
This PR fixes #19818 by adding a more descriptive message to the mask mismatch exception in batch normalization.

Before:
![image](https://github.com/keras-team/keras/assets/60848863/5303a7b4-618b-484e-832f-e8b319877af5)


After (Highlighted for clarity):
![image](https://github.com/keras-team/keras/assets/60848863/c88256f4-2b02-445b-bc7e-744e051ff958)


Patch:
diff --git a/keras/api/_tf_keras/keras/ops/__init__.py b/keras/api/_tf_keras/keras/ops/__init__.py
index 7d9b06886..d287ddd6c 100644
--- a/keras/api/_tf_keras/keras/ops/__init__.py
+++ b/keras/api/_tf_keras/keras/ops/__init__.py
@@ -22,6 +22,7 @@ from keras.src.ops.core import shape
 from keras.src.ops.core import slice
 from keras.src.ops.core import slice_update
 from keras.src.ops.core import stop_gradient
+from keras.src.ops.core import switch
 from keras.src.ops.core import unstack
 from keras.src.ops.core import vectorized_map
 from keras.src.ops.core import while_loop
diff --git a/keras/api/ops/__init__.py b/keras/api/ops/__init__.py
index 7d9b06886..d287ddd6c 100644
--- a/keras/api/ops/__init__.py
+++ b/keras/api/ops/__init__.py
@@ -22,6 +22,7 @@ from keras.src.ops.core import shape
 from keras.src.ops.core import slice
 from keras.src.ops.core import slice_update
 from keras.src.ops.core import stop_gradient
+from keras.src.ops.core import switch
 from keras.src.ops.core import unstack
 from keras.src.ops.core import vectorized_map
 from keras.src.ops.core import while_loop
diff --git a/keras/src/backend/jax/core.py b/keras/src/backend/jax/core.py
index d4682ffff..fba75be1f 100644
--- a/keras/src/backend/jax/core.py
+++ b/keras/src/backend/jax/core.py
@@ -287,6 +287,10 @@ def slice_update(inputs, start_indices, updates):
     return jax.lax.dynamic_update_slice(inputs, updates, start_indices)
 
 
+def switch(index, branches, *operands):
+    return jax.lax.switch(index, branches, *operands)
+
+
 def while_loop(
     cond,
     body,
diff --git a/keras/src/backend/numpy/core.py b/keras/src/backend/numpy/core.py
index 25fb1ea07..31700428d 100644
--- a/keras/src/backend/numpy/core.py
+++ b/keras/src/backend/numpy/core.py
@@ -241,6 +241,12 @@ def slice_update(inputs, start_indices, updates):
     return inputs
 
 
+def switch(index, branches, *operands):
+    index = convert_to_tensor(index, "int32")
+    index = np.clip(index, 0, len(branches) - 1)
+    return branches[index](*operands)
+
+
 def while_loop(
     cond,
     body,
diff --git a/keras/src/backend/tensorflow/core.py b/keras/src/backend/tensorflow/core.py
index 5c44b47e0..82a797d3e 100644
--- a/keras/src/backend/tensorflow/core.py
+++ b/keras/src/backend/tensorflow/core.py
@@ -355,6 +355,19 @@ def slice_update(inputs, start_indices, updates):
     return dynamic_update_slice(inputs, updates, start_indices)
 
 
+def switch(index, branches, *operands):
+    index = convert_to_tensor(index, "int32")
+    index = tf.clip_by_value(index, 0, len(branches) - 1)
+
+    # Workaround to deal with python closures. More details:
+    # https://github.com/tensorflow/tensorflow/issues/8776#issuecomment-311383887
+    def gen_fn(i):
+        return lambda: branches[i](*operands)
+
+    branch_fns = [gen_fn(i) for i in range(len(branches))]
+    return tf.switch_case(index, branch_fns)
+
+
 def while_loop(
     cond,
     body,
diff --git a/keras/src/backend/torch/core.py b/keras/src/backend/torch/core.py
index 3fb14a75b..2f2871b24 100644
--- a/keras/src/backend/torch/core.py
+++ b/keras/src/backend/torch/core.py
@@ -445,6 +445,12 @@ def slice_update(inputs, start_indices, updates):
     return outputs
 
 
+def switch(index, branches, *operands):
+    index = convert_to_tensor(index, "int32")
+    index = torch.clamp(index, 0, len(branches) - 1)
+    return branches[index](*operands)
+
+
 def while_loop(
     cond,
     body,
diff --git a/keras/src/layers/normalization/batch_normalization.py b/keras/src/layers/normalization/batch_normalization.py
index ecbd1a453..642dbd4b1 100644
--- a/keras/src/layers/normalization/batch_normalization.py
+++ b/keras/src/layers/normalization/batch_normalization.py
@@ -214,6 +214,15 @@ class BatchNormalization(Layer):
         return input_shape
 
     def call(self, inputs, training=None, mask=None):
+        # Check if the mask has one less dimension than the inputs.
+        if mask is not None:
+            if len(mask.shape) != len(inputs.shape) - 1:
+                # Raise a value error
+                raise ValueError(
+                    "The mask provided should be one dimension less than the inputs. "
+                    f"Received: mask.shape={mask.shape}, inputs.shape={inputs.shape}"
+                )
+
         input_dtype = standardize_dtype(inputs.dtype)
         if input_dtype in ("float16", "bfloat16"):
             # BN is prone to overflowing for float16/bfloat16 inputs, so we opt
diff --git a/keras/src/metrics/metric.py b/keras/src/metrics/metric.py
index 91e50dab8..0a3891473 100644
--- a/keras/src/metrics/metric.py
+++ b/keras/src/metrics/metric.py
@@ -209,9 +209,9 @@ class Metric(KerasSaveable):
 
     @property
     def variables(self):
-        variables = self._variables[:]
+        variables = list(self._variables)
         for metric in self._metrics:
-            variables.extend(metric._variables)
+            variables.extend(metric.variables)
         return variables
 
     def __call__(self, *args, **kwargs):
diff --git a/keras/src/ops/core.py b/keras/src/ops/core.py
index 3b5841290..5d7fdf2ff 100644
--- a/keras/src/ops/core.py
+++ b/keras/src/ops/core.py
@@ -118,7 +118,7 @@ def scan(f, init, xs=None, length=None, reverse=False, unroll=1):
     >>> sum_fn = lambda c, x: (c + x, c + x)
     >>> init = keras.ops.array(0)
     >>> xs = keras.ops.array([1, 2, 3, 4, 5])
-    >>> carry, result = ops.scan(sum_fn, init, xs)
+    >>> carry, result = keras.ops.scan(sum_fn, init, xs)
     >>> carry
     15
     >>> result
@@ -315,6 +315,57 @@ def slice_update(inputs, start_indices, updates):
     return backend.core.slice_update(inputs, start_indices, updates)
 
 
+class Switch(Operation):
+    def call(self, index, branches, *operands):
+        return backend.core.switch(index, branches, *operands)
+
+    def compute_output_spec(self, index, branches, *operands):
+        # We use first branch for output_spec
+        spec = backend.compute_output_spec(branches[0], *operands)
+        return spec
+
+
+@keras_export("keras.ops.switch")
+def switch(index, branches, *operands):
+    """Apply exactly one of the `branches` given by `index`.
+
+    If `index` is out of bounds, it is clamped to within bounds.
+
+    The semantics of `switch` are given roughly by this Python implementation:
+
+    ```python
+    def switch(index, branches, *operands):
+        index = clamp(0, index, len(branches) - 1)
+        return branches[index](*operands)
+    ```
+
+    Args:
+        index: An integer scalar indicating which branch function to apply.
+        branches: A sequence of functions to be applied based on `index`.
+        operands: Inputs to whichever branch is applied.
+
+    Returns:
+        The outputs of `branch(*operands)` for the branch that was selected
+        based on `index`.
+
+    Examples:
+
+    >>> add_fn = lambda x, y: x + y
+    >>> substract_fn = lambda x, y: x - y
+    >>> x = keras.ops.array(2.0)
+    >>> y = keras.ops.array(0.5)
+    >>> branches = [add_fn, substract_fn]
+    >>> keras.ops.switch(0, branches, x, y)
+    2.5
+
+    >>> keras.ops.switch(1, branches, x, y)
+    1.5
+    """
+    if any_symbolic_tensors(operands):
+        return Switch().symbolic_call(index, branches, *operands)
+    return backend.core.switch(index, branches, *operands)
+
+
 class WhileLoop(Operation):
     def __init__(self, cond, body, maximum_iterations):
         super().__init__()
diff --git a/keras/src/ops/numpy.py b/keras/src/ops/numpy.py
index e8289ee2e..69dbe2307 100644
--- a/keras/src/ops/numpy.py
+++ b/keras/src/ops/numpy.py
@@ -4345,26 +4345,44 @@ class Repeat(Operation):
 
     def compute_output_spec(self, x):
         x_shape = list(x.shape)
+        repeats = self.repeats
+        if isinstance(repeats, int):
+            repeats = [repeats]
+        repeats_size = len(repeats)
+        broadcast = repeats_size == 1
+
         if self.axis is None:
             if None in x_shape:
                 return KerasTensor([None], dtype=x.dtype)
 
             x_flatten_size = int(np.prod(x_shape))
-            if isinstance(self.repeats, int):
-                output_shape = [x_flatten_size * self.repeats]
+            if broadcast:
+                output_shape = [x_flatten_size * repeats[0]]
+            elif repeats_size != x_flatten_size:
+                raise ValueError(
+                    "Size of `repeats` and "
+                    "dimensions of `x` after flattening should be compatible. "
+                    f"Received: {repeats_size} and {x_flatten_size}"
+                )
             else:
-                output_shape = [int(np.sum(self.repeats))]
+                output_shape = [int(np.sum(repeats))]
             return KerasTensor(output_shape, dtype=x.dtype)
 
         size_on_ax = x_shape[self.axis]
+        if size_on_ax is None:
+            return KerasTensor(x_shape, dtype=x.dtype)
+
         output_shape = x_shape
-        if isinstance(self.repeats, int):
-            if size_on_ax is None:
-                output_shape[self.axis] = None
-            else:
-                output_shape[self.axis] = size_on_ax * self.repeats
+        if broadcast:
+            output_shape[self.axis] = size_on_ax * repeats[0]
+        elif size_on_ax != repeats_size:
+            raise ValueError(
+                "Size of `repeats` and "
+                f"dimensions of `axis {self.axis} of x` should be compatible. "
+                f"Received: {repeats_size} and {x_shape}"
+            )
         else:
-            output_shape[self.axis] = int(np.sum(self.repeats))
+            output_shape[self.axis] = int(np.sum(repeats))
         return KerasTensor(output_shape, dtype=x.dtype)
 
 
diff --git a/keras/src/utils/file_utils.py b/keras/src/utils/file_utils.py
index e625a9f13..c3a42d54f 100644
--- a/keras/src/utils/file_utils.py
+++ b/keras/src/utils/file_utils.py
@@ -270,9 +270,9 @@ def get_file(
                 self.finished = False
 
             def __call__(self, block_num, block_size, total_size):
+                if total_size == -1:
+                    total_size = None
                 if not self.progbar:
-                    if total_size == -1:
-                        total_size = None
                     self.progbar = Progbar(total_size)
                 current = block_num * block_size
 
diff --git a/requirements-common.txt b/requirements-common.txt
index 29090e9c9..54e6d45f7 100644
--- a/requirements-common.txt
+++ b/requirements-common.txt
@@ -19,5 +19,4 @@ optree
 pytest-cov
 packaging
 # for tree_test.py
-parameterized
-dm_tree
\ No newline at end of file
+dm_tree


Test patch:
diff --git a/keras/src/metrics/metric_test.py b/keras/src/metrics/metric_test.py
index 90bd1a9b9..0d9635a25 100644
--- a/keras/src/metrics/metric_test.py
+++ b/keras/src/metrics/metric_test.py
@@ -161,6 +161,12 @@ class MetricTest(testing.TestCase):
         }
         self.assertEqual(len(metric.variables), 6)
 
+        # Two levels deep
+        metric = ExampleMetric(name="mse")
+        metric.submetric = ExampleMetric(name="submse")
+        metric.submetric.submetric = ExampleMetric(name="subsubmse")
+        self.assertEqual(len(metric.variables), 6)
+
     def test_serialization(self):
         self.run_class_serialization_test(
             ExampleMetric(name="mse"),
diff --git a/keras/src/ops/core_test.py b/keras/src/ops/core_test.py
index 4309bc091..91b788e75 100644
--- a/keras/src/ops/core_test.py
+++ b/keras/src/ops/core_test.py
@@ -73,6 +73,16 @@ class CoreOpsStaticShapeTest(testing.TestCase):
             core.slice_update(inputs, start_indices, updates).shape, (4, 4, 4)
         )
 
+    def test_switch(self):
+        def fn(x, y):
+            return x[:, 0], y[0, :]
+
+        index = KerasTensor(())
+        x = KerasTensor((5, 2))
+        y = KerasTensor((5, 2))
+        self.assertEqual(core.switch(index, [fn], x, y)[0].shape, (5,))
+        self.assertEqual(core.switch(index, [fn], x, y)[1].shape, (2,))
+
     def test_fori_loop(self):
         def body_fun(i, x):
             return x + i
@@ -303,6 +313,23 @@ class CoreOpsCorrectnessTest(testing.TestCase, parameterized.TestCase):
         outputs = core.slice_update(inputs, start_indices, updates)
         self.assertAllClose(outputs[1:3, 1:3, 2:4, 2:4], np.zeros([2, 2, 2, 2]))
 
+    def test_switch(self):
+        def fn1(x, y):
+            return x + y
+
+        def fn2(x, y):
+            return x - y
+
+        x = np.random.rand(2, 3, 4).astype("float32")
+        y = np.random.rand(2, 3, 4).astype("float32")
+        branches = [fn1, fn2]
+        self.assertAllClose(core.switch(0, branches, x, y), x + y)
+        self.assertAllClose(core.switch(1, branches, x, y), x - y)
+
+        # Test out-of-bound index
+        self.assertAllClose(core.switch(-100, branches, x, y), x + y)
+        self.assertAllClose(core.switch(100, branches, x, y), x - y)
+
     @parameterized.named_parameters(
         [
             {
@@ -801,6 +828,25 @@ class CoreOpsCallsTests(testing.TestCase):
         expected_output = np.array([[1, 2, 3], [4, 10, 11], [7, 12, 13]])
         self.assertAllClose(core.convert_to_numpy(result), expected_output)
 
+    def test_switch_basic_call(self):
+        def fn1(x, y):
+            return x + y
+
+        def fn2(x, y):
+            return x - y
+
+        x = np.random.rand(2, 3, 4).astype("float32")
+        y = np.random.rand(2, 3, 4).astype("float32")
+        branches = [fn1, fn2]
+        switch_op = core.Switch()
+        index = 0
+        outputs = switch_op.call(index, branches, x, y)
+        self.assertAllClose(outputs, x + y)
+
+        index = 1
+        outputs = switch_op.call(index, branches, x, y)
+        self.assertAllClose(outputs, x - y)
+
     def test_while_loop_basic_functionality(self):
         # Loop condition: continue if i < 5
         def cond(i):
diff --git a/keras/src/ops/numpy_test.py b/keras/src/ops/numpy_test.py
index e257cde4a..ab730608f 100644
--- a/keras/src/ops/numpy_test.py
+++ b/keras/src/ops/numpy_test.py
@@ -1364,7 +1364,7 @@ class NumpyOneInputOpsDynamicShapeTest(testing.TestCase):
         x = KerasTensor((None, 3))
         self.assertEqual(knp.repeat(x, 2).shape, (None,))
         self.assertEqual(knp.repeat(x, 3, axis=1).shape, (None, 9))
-        self.assertEqual(knp.repeat(x, [1, 2], axis=0).shape, (3, 3))
+        self.assertEqual(knp.repeat(x, [1, 2], axis=0).shape, (None, 3))
         self.assertEqual(knp.repeat(x, 2, axis=0).shape, (None, 3))
 
     def test_reshape(self):
@@ -1875,9 +1875,15 @@ class NumpyOneInputOpsStaticShapeTest(testing.TestCase):
     def test_repeat(self):
         x = KerasTensor((2, 3))
         self.assertEqual(knp.repeat(x, 2).shape, (12,))
+        self.assertEqual(knp.repeat(x, [2]).shape, (12,))
         self.assertEqual(knp.repeat(x, 3, axis=1).shape, (2, 9))
         self.assertEqual(knp.repeat(x, [1, 2], axis=0).shape, (3, 3))
 
+        with self.assertRaises(ValueError):
+            knp.repeat(x, [1, 1])
+        with self.assertRaises(ValueError):
+            knp.repeat(x, [1, 1, 1], axis=0)
+
     def test_reshape(self):
         x = KerasTensor((2, 3))
         self.assertEqual(knp.reshape(x, (3, 2)).shape, (3, 2))
@@ -3902,6 +3908,10 @@ class NumpyOneInputOpsCorrectnessTest(testing.TestCase, parameterized.TestCase):
     def test_repeat(self):
         x = np.array([[1, 2], [3, 4]])
         self.assertAllClose(knp.repeat(x, 2), np.repeat(x, 2))
+        self.assertAllClose(
+            knp.Repeat(np.array([2]))(x),
+            np.repeat(x, np.array([2])),
+        )
         self.assertAllClose(knp.repeat(x, 3, axis=1), np.repeat(x, 3, axis=1))
         self.assertAllClose(
             knp.repeat(x, np.array([1, 2]), axis=-1),
diff --git a/keras/src/tree/tree_test.py b/keras/src/tree/tree_test.py
index 6859ac4a6..a560d5009 100644
--- a/keras/src/tree/tree_test.py
+++ b/keras/src/tree/tree_test.py
@@ -1,7 +1,7 @@
 import collections
 
 import numpy as np
-import parameterized
+from absl.testing import parameterized
 
 from keras.src import ops
 from keras.src import testing
@@ -13,97 +13,89 @@ STRUCTURE2 = ((("foo1", "foo2"), "foo3"), "foo4", ("foo5", "foo6"))
 STRUCTURE_DIFFERENT_NUM_ELEMENTS = ("spam", "eggs")
 STRUCTURE_DIFFERENT_NESTING = (((1, 2), 3), 4, 5, (6,))
 
+TEST_CASES = []
+if dmtree.available:
+    from keras.src.tree import dmtree_impl
 
-@parameterized.parameterized_class(
-    ("tree_module",),
-    [
-        (dmtree,),
-        (optree,),
-    ],
-)
-class TreeTest(testing.TestCase):
-    tree_impl = None
-    tree_api_name = ""
+    TEST_CASES += [
+        {
+            "testcase_name": "dmtree",
+            "tree_impl": dmtree_impl,
+            "is_optree": False,
+        }
+    ]
+if optree.available:
+    from keras.src.tree import optree_impl
 
-    def setUp(self):
-        super().setUp()
-        if self.tree_module.available:
-            if self.tree_module == dmtree:
-                from keras.src.tree import dmtree_impl as tree_impl
+    TEST_CASES += [
+        {
+            "testcase_name": "optree",
+            "tree_impl": optree_impl,
+            "is_optree": True,
+        },
+    ]
 
-                self.tree_impl = tree_impl
-                self.tree_api_name = "dmtree"
-            elif self.tree_module == optree:
-                from keras.src.tree import optree_impl as tree_impl
 
-                self.tree_impl = tree_impl
-                self.tree_api_name = "optree"
-            else:
-                raise ValueError(f"Unrecognized module {self.tree_module}.")
-        else:
-            self.skipTest(f"Skip since {self.tree_module} is not available.")
+@parameterized.named_parameters(TEST_CASES)
+class TreeTest(testing.TestCase):
 
-    def test_is_nested(self):
-        self.assertFalse(self.tree_impl.is_nested("1234"))
-        self.assertFalse(self.tree_impl.is_nested(b"1234"))
-        self.assertFalse(self.tree_impl.is_nested(bytearray("1234", "ascii")))
-        self.assertTrue(self.tree_impl.is_nested([1, 3, [4, 5]]))
-        self.assertTrue(self.tree_impl.is_nested(((7, 8), (5, 6))))
-        self.assertTrue(self.tree_impl.is_nested([]))
-        self.assertTrue(self.tree_impl.is_nested({"a": 1, "b": 2}))
-        self.assertFalse(self.tree_impl.is_nested(set([1, 2])))
+    def test_is_nested(self, tree_impl, is_optree):
+        self.assertFalse(tree_impl.is_nested("1234"))
+        self.assertFalse(tree_impl.is_nested(b"1234"))
+        self.assertFalse(tree_impl.is_nested(bytearray("1234", "ascii")))
+        self.assertTrue(tree_impl.is_nested([1, 3, [4, 5]]))
+        self.assertTrue(tree_impl.is_nested(((7, 8), (5, 6))))
+        self.assertTrue(tree_impl.is_nested([]))
+        self.assertTrue(tree_impl.is_nested({"a": 1, "b": 2}))
+        self.assertFalse(tree_impl.is_nested(set([1, 2])))
         ones = np.ones([2, 3])
-        self.assertFalse(self.tree_impl.is_nested(ones))
-        self.assertFalse(self.tree_impl.is_nested(np.tanh(ones)))
-        self.assertFalse(self.tree_impl.is_nested(np.ones((4, 5))))
+        self.assertFalse(tree_impl.is_nested(ones))
+        self.assertFalse(tree_impl.is_nested(np.tanh(ones)))
+        self.assertFalse(tree_impl.is_nested(np.ones((4, 5))))
 
-    def test_flatten(self):
+    def test_flatten(self, tree_impl, is_optree):
         structure = ((3, 4), 5, (6, 7, (9, 10), 8))
         flat = ["a", "b", "c", "d", "e", "f", "g", "h"]
 
         self.assertEqual(
-            self.tree_impl.flatten(structure), [3, 4, 5, 6, 7, 9, 10, 8]
+            tree_impl.flatten(structure), [3, 4, 5, 6, 7, 9, 10, 8]
         )
         point = collections.namedtuple("Point", ["x", "y"])
         structure = (point(x=4, y=2), ((point(x=1, y=0),),))
         flat = [4, 2, 1, 0]
-        self.assertEqual(self.tree_impl.flatten(structure), flat)
+        self.assertEqual(tree_impl.flatten(structure), flat)
 
-        self.assertEqual([5], self.tree_impl.flatten(5))
-        self.assertEqual([np.array([5])], self.tree_impl.flatten(np.array([5])))
+        self.assertEqual([5], tree_impl.flatten(5))
+        self.assertEqual([np.array([5])], tree_impl.flatten(np.array([5])))
 
-    def test_flatten_dict_order(self):
+    def test_flatten_dict_order(self, tree_impl, is_optree):
         ordered = collections.OrderedDict(
             [("d", 3), ("b", 1), ("a", 0), ("c", 2)]
         )
         plain = {"d": 3, "b": 1, "a": 0, "c": 2}
-        ordered_flat = self.tree_impl.flatten(ordered)
-        plain_flat = self.tree_impl.flatten(plain)
+        ordered_flat = tree_impl.flatten(ordered)
+        plain_flat = tree_impl.flatten(plain)
         # dmtree does not respect the ordered dict.
-        if self.tree_api_name == "optree":
+        if is_optree:
             self.assertEqual([3, 1, 0, 2], ordered_flat)
         else:
             self.assertEqual([0, 1, 2, 3], ordered_flat)
         self.assertEqual([0, 1, 2, 3], plain_flat)
 
-    def test_map_structure(self):
+    def test_map_structure(self, tree_impl, is_optree):
         assertion_message = (
             "have the same structure"
-            if self.tree_api_name == "optree"
+            if is_optree
             else "have the same nested structure"
         )
-        assertion_type_error = (
-            ValueError if self.tree_api_name == "optree" else TypeError
-        )
+        assertion_type_error = ValueError if is_optree else TypeError
         structure2 = (((7, 8), 9), 10, (11, 12))
-        structure1_plus1 = self.tree_impl.map_structure(
-            lambda x: x + 1, STRUCTURE1
-        )
-        self.tree_impl.assert_same_structure(STRUCTURE1, structure1_plus1)
+        structure1_plus1 = tree_impl.map_structure(lambda x: x + 1, STRUCTURE1)
+        tree_impl.assert_same_structure(STRUCTURE1, structure1_plus1)
         self.assertAllEqual(
-            [2, 3, 4, 5, 6, 7], self.tree_impl.flatten(structure1_plus1)
+            [2, 3, 4, 5, 6, 7], tree_impl.flatten(structure1_plus1)
         )
-        structure1_plus_structure2 = self.tree_impl.map_structure(
+        structure1_plus_structure2 = tree_impl.map_structure(
             lambda x, y: x + y, STRUCTURE1, structure2
         )
         self.assertEqual(
@@ -111,57 +103,49 @@ class TreeTest(testing.TestCase):
             structure1_plus_structure2,
         )
 
-        self.assertEqual(3, self.tree_impl.map_structure(lambda x: x - 1, 4))
+        self.assertEqual(3, tree_impl.map_structure(lambda x: x - 1, 4))
 
-        self.assertEqual(
-            7, self.tree_impl.map_structure(lambda x, y: x + y, 3, 4)
-        )
+        self.assertEqual(7, tree_impl.map_structure(lambda x, y: x + y, 3, 4))
 
         # Empty structures
-        self.assertEqual((), self.tree_impl.map_structure(lambda x: x + 1, ()))
-        self.assertEqual([], self.tree_impl.map_structure(lambda x: x + 1, []))
-        self.assertEqual({}, self.tree_impl.map_structure(lambda x: x + 1, {}))
+        self.assertEqual((), tree_impl.map_structure(lambda x: x + 1, ()))
+        self.assertEqual([], tree_impl.map_structure(lambda x: x + 1, []))
+        self.assertEqual({}, tree_impl.map_structure(lambda x: x + 1, {}))
         empty_nt = collections.namedtuple("empty_nt", "")
         self.assertEqual(
             empty_nt(),
-            self.tree_impl.map_structure(lambda x: x + 1, empty_nt()),
+            tree_impl.map_structure(lambda x: x + 1, empty_nt()),
         )
 
         # This is checking actual equality of types, empty list != empty tuple
-        self.assertNotEqual(
-            (), self.tree_impl.map_structure(lambda x: x + 1, [])
-        )
+        self.assertNotEqual((), tree_impl.map_structure(lambda x: x + 1, []))
 
         with self.assertRaisesRegex(TypeError, "callable"):
-            self.tree_impl.map_structure("bad", structure1_plus1)
+            tree_impl.map_structure("bad", structure1_plus1)
         with self.assertRaisesRegex(ValueError, "at least one structure"):
-            self.tree_impl.map_structure(lambda x: x)
+            tree_impl.map_structure(lambda x: x)
         with self.assertRaisesRegex(ValueError, assertion_message):
-            self.tree_impl.map_structure(lambda x, y: None, (3, 4), (3, 4, 5))
+            tree_impl.map_structure(lambda x, y: None, (3, 4), (3, 4, 5))
         with self.assertRaisesRegex(ValueError, assertion_message):
-            self.tree_impl.map_structure(lambda x, y: None, 3, (3,))
+            tree_impl.map_structure(lambda x, y: None, 3, (3,))
         with self.assertRaisesRegex(assertion_type_error, assertion_message):
-            self.tree_impl.map_structure(
-                lambda x, y: None, ((3, 4), 5), [(3, 4), 5]
-            )
+            tree_impl.map_structure(lambda x, y: None, ((3, 4), 5), [(3, 4), 5])
         with self.assertRaisesRegex(ValueError, assertion_message):
-            self.tree_impl.map_structure(
-                lambda x, y: None, ((3, 4), 5), (3, (4, 5))
-            )
+            tree_impl.map_structure(lambda x, y: None, ((3, 4), 5), (3, (4, 5)))
 
         structure1_list = [[[1, 2], 3], 4, [5, 6]]
         with self.assertRaisesRegex(assertion_type_error, assertion_message):
-            self.tree_impl.map_structure(
+            tree_impl.map_structure(
                 lambda x, y: None, STRUCTURE1, structure1_list
             )
 
-    def test_map_structure_up_to(self):
+    def test_map_structure_up_to(self, tree_impl, is_optree):
         # Named tuples.
         ab_tuple = collections.namedtuple("ab_tuple", "a, b")
         op_tuple = collections.namedtuple("op_tuple", "add, mul")
         inp_val = ab_tuple(a=2, b=3)
         inp_ops = ab_tuple(a=op_tuple(add=1, mul=2), b=op_tuple(add=2, mul=3))
-        out = self.tree_impl.map_structure_up_to(
+        out = tree_impl.map_structure_up_to(
             inp_val,
             lambda val, ops: (val + ops.add) * ops.mul,
             inp_val,
@@ -173,7 +157,7 @@ class TreeTest(testing.TestCase):
         # Lists.
         data_list = [[2, 4, 6, 8], [[1, 3, 5, 7, 9], [3, 5, 7]]]
         name_list = ["evens", ["odds", "primes"]]
-        out = self.tree_impl.map_structure_up_to(
+        out = tree_impl.map_structure_up_to(
             name_list,
             lambda name, sec: "first_{}_{}".format(len(sec), name),
             name_list,
@@ -183,77 +167,75 @@ class TreeTest(testing.TestCase):
             out, ["first_4_evens", ["first_5_odds", "first_3_primes"]]
         )
 
-    def test_assert_same_structure(self):
+    def test_assert_same_structure(self, tree_impl, is_optree):
         assertion_message = (
             "have the same structure"
-            if self.tree_api_name == "optree"
+            if is_optree
             else "have the same nested structure"
         )
-        assertion_type_error = (
-            ValueError if self.tree_api_name == "optree" else TypeError
-        )
-        self.tree_impl.assert_same_structure(
+        assertion_type_error = ValueError if is_optree else TypeError
+        tree_impl.assert_same_structure(
             STRUCTURE1, STRUCTURE2, check_types=False
         )
-        self.tree_impl.assert_same_structure("abc", 1.0, check_types=False)
-        self.tree_impl.assert_same_structure(b"abc", 1.0, check_types=False)
-        self.tree_impl.assert_same_structure("abc", 1.0, check_types=False)
-        self.tree_impl.assert_same_structure(
+        tree_impl.assert_same_structure("abc", 1.0, check_types=False)
+        tree_impl.assert_same_structure(b"abc", 1.0, check_types=False)
+        tree_impl.assert_same_structure("abc", 1.0, check_types=False)
+        tree_impl.assert_same_structure(
             bytearray("abc", "ascii"), 1.0, check_types=False
         )
-        self.tree_impl.assert_same_structure(
+        tree_impl.assert_same_structure(
             "abc", np.array([0, 1]), check_types=False
         )
 
         with self.assertRaisesRegex(ValueError, assertion_message):
-            self.tree_impl.assert_same_structure(
+            tree_impl.assert_same_structure(
                 STRUCTURE1, STRUCTURE_DIFFERENT_NUM_ELEMENTS
             )
         with self.assertRaisesRegex(ValueError, assertion_message):
-            self.tree_impl.assert_same_structure([0, 1], np.array([0, 1]))
+            tree_impl.assert_same_structure([0, 1], np.array([0, 1]))
         with self.assertRaisesRegex(ValueError, assertion_message):
-            self.tree_impl.assert_same_structure(0, [0, 1])
+            tree_impl.assert_same_structure(0, [0, 1])
         with self.assertRaisesRegex(assertion_type_error, assertion_message):
-            self.tree_impl.assert_same_structure((0, 1), [0, 1])
+            tree_impl.assert_same_structure((0, 1), [0, 1])
         with self.assertRaisesRegex(ValueError, assertion_message):
-            self.tree_impl.assert_same_structure(
+            tree_impl.assert_same_structure(
                 STRUCTURE1, STRUCTURE_DIFFERENT_NESTING
             )
         with self.assertRaisesRegex(ValueError, assertion_message):
-            self.tree_impl.assert_same_structure([[3], 4], [3, [4]])
+            tree_impl.assert_same_structure([[3], 4], [3, [4]])
         with self.assertRaisesRegex(ValueError, assertion_message):
-            self.tree_impl.assert_same_structure({"a": 1}, {"b": 1})
+            tree_impl.assert_same_structure({"a": 1}, {"b": 1})
         structure1_list = [[[1, 2], 3], 4, [5, 6]]
         with self.assertRaisesRegex(assertion_type_error, assertion_message):
-            self.tree_impl.assert_same_structure(STRUCTURE1, structure1_list)
-        self.tree_impl.assert_same_structure(
+            tree_impl.assert_same_structure(STRUCTURE1, structure1_list)
+        tree_impl.assert_same_structure(
             STRUCTURE1, STRUCTURE2, check_types=False
         )
         # dm-tree treat list and tuple only on type mismatch, but optree treat
         # them as structure mismatch.
-        if self.tree_api_name == "optree":
+        if is_optree:
             with self.assertRaisesRegex(
                 assertion_type_error, assertion_message
             ):
-                self.tree_impl.assert_same_structure(
+                tree_impl.assert_same_structure(
                     STRUCTURE1, structure1_list, check_types=False
                 )
         else:
-            self.tree_impl.assert_same_structure(
+            tree_impl.assert_same_structure(
                 STRUCTURE1, structure1_list, check_types=False
             )
 
-    def test_pack_sequence_as(self):
+    def test_pack_sequence_as(self, tree_impl, is_optree):
         structure = {"key3": "", "key1": "", "key2": ""}
         flat_sequence = ["value1", "value2", "value3"]
         self.assertEqual(
-            self.tree_impl.pack_sequence_as(structure, flat_sequence),
+            tree_impl.pack_sequence_as(structure, flat_sequence),
             {"key3": "value3", "key1": "value1", "key2": "value2"},
         )
         structure = (("a", "b"), ("c", "d", "e"), "f")
         flat_sequence = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
         self.assertEqual(
-            self.tree_impl.pack_sequence_as(structure, flat_sequence),
+            tree_impl.pack_sequence_as(structure, flat_sequence),
             ((1.0, 2.0), (3.0, 4.0, 5.0), 6.0),
         )
         structure = {
@@ -262,7 +244,7 @@ class TreeTest(testing.TestCase):
         }
         flat_sequence = ["val2", "val1", 3.0, 1.0, 2.0]
         self.assertEqual(
-            self.tree_impl.pack_sequence_as(structure, flat_sequence),
+            tree_impl.pack_sequence_as(structure, flat_sequence),
             {
                 "key3": {"c": (1.0, 2.0), "a": 3.0},
                 "key1": {"e": "val1", "d": "val2"},
@@ -271,43 +253,41 @@ class TreeTest(testing.TestCase):
         structure = ["a"]
         flat_sequence = [np.array([[1, 2], [3, 4]])]
         self.assertAllClose(
-            self.tree_impl.pack_sequence_as(structure, flat_sequence),
+            tree_impl.pack_sequence_as(structure, flat_sequence),
             [np.array([[1, 2], [3, 4]])],
         )
         structure = ["a"]
         flat_sequence = [ops.ones([2, 2])]
         self.assertAllClose(
-            self.tree_impl.pack_sequence_as(structure, flat_sequence),
+            tree_impl.pack_sequence_as(structure, flat_sequence),
             [ops.ones([2, 2])],
         )
 
         with self.assertRaisesRegex(TypeError, "Attempted to pack value:"):
             structure = ["a"]
             flat_sequence = 1
-            self.tree_impl.pack_sequence_as(structure, flat_sequence)
+            tree_impl.pack_sequence_as(structure, flat_sequence)
         with self.assertRaisesRegex(ValueError, "The target structure is of"):
             structure = "a"
             flat_sequence = [1, 2]
-            self.tree_impl.pack_sequence_as(structure, flat_sequence)
+            tree_impl.pack_sequence_as(structure, flat_sequence)
 
-    def test_lists_to_tuples(self):
+    def test_lists_to_tuples(self, tree_impl, is_optree):
         structure = [1, 2, 3]
-        self.assertEqual(self.tree_impl.lists_to_tuples(structure), (1, 2, 3))
+        self.assertEqual(tree_impl.lists_to_tuples(structure), (1, 2, 3))
         structure = [[1], [2, 3]]
-        self.assertEqual(
-            self.tree_impl.lists_to_tuples(structure), ((1,), (2, 3))
-        )
+        self.assertEqual(tree_impl.lists_to_tuples(structure), ((1,), (2, 3)))
         structure = [[1], [2, [3]]]
         self.assertEqual(
-            self.tree_impl.lists_to_tuples(structure), ((1,), (2, (3,)))
+            tree_impl.lists_to_tuples(structure), ((1,), (2, (3,)))
         )
 
-    def test_traverse(self):
+    def test_traverse(self, tree_impl, is_optree):
         # Lists to tuples
         structure = [(1, 2), [3], {"a": [4]}]
         self.assertEqual(
             ((1, 2), (3,), {"a": (4,)}),
-            self.tree_impl.traverse(
+            tree_impl.traverse(
                 lambda x: tuple(x) if isinstance(x, list) else x,
                 structure,
                 top_down=False,
@@ -321,7 +301,7 @@ class TreeTest(testing.TestCase):
             visited.append(x)
             return "X" if isinstance(x, tuple) and len(x) > 2 else None
 
-        output = self.tree_impl.traverse(visit, structure)
+        output = tree_impl.traverse(visit, structure)
         self.assertEqual([(1, [2]), [3, "X"]], output)
         self.assertEqual(
             [