Method test_stack_call

tests/nn/layer_stack_test.py:46–109

(self)

class LayerStackTest(absltest.TestCase):

  def test_stack_call(self):
    the_layer = pz.nn.LayerStack(
        stacked_sublayers=MyTestLayer(
            params=(
                pz.Parameter(
                    value=pz.nx.zeros({"stack": 3, "foo": 4}), label="param1"
                ),
                pz.Parameter(value=pz.nx.zeros({"bar": 5}), label="param2"),
            ),
            expected_param_shapes=({"foo": 4}, {"bar": 5}),
            states=(
                pz.StateVariable(
                    value=0,
                    label="state1",
                    metadata={
                        "layerstack_axes": {
                            "stack": pz.nn.LayerStackVarBehavior.SHARED
                        }
                    },
                ),
                pz.StateVariable(
                    value=pz.nx.zeros({"foo": 3}),
                    label="state2",
                    metadata={
                        "layerstack_axes": {
                            "stack": pz.nn.LayerStackVarBehavior.SHARED
                        }
                    },
                ),
                pz.StateVariable(
                    value=pz.nx.arange("stack", 3),
                    label="state3",
                    metadata={
                        "layerstack_axes": {
                            "stack": pz.nn.LayerStackVarBehavior.PER_LAYER
                        }
                    },
                ),
            ),
            stuff={
                "value": 100.0,
                "named": pz.nx.arange("bar", 3),
                "named_stacked": pz.nx.arange("stack", 3),
            },
        ),
        stack_axis="stack",
        stack_axis_size=3,
    )
    counter = pz.StateVariable(0)
    result = the_layer(10, shared_counter=counter)

    self.assertEqual(result, 13.0)
    self.assertEqual(counter.value, 3)
    chex.assert_trees_all_equal(
        the_layer.stacked_sublayers.states[0].value, 3.0
    )
    chex.assert_trees_all_equal(
        the_layer.stacked_sublayers.states[1].value.canonicalize(),