Method test_layer_stack_build

tests/nn/layer_stack_test.py:111–179

(self)

    )

  def test_layer_stack_build(self):
    def builder(init_base_rng, some_value):
      return MyTestLayer(
          params=(
              pz.nn.make_parameter(
                  "foo",
                  init_base_rng,
                  lambda k: pz.nx.wrap(
                      jax.random.uniform(k, shape=(4,)), "vals"
                  ),
              ),
          ),
          expected_param_shapes=None,
          states=(pz.StateVariable(value=pz.nx.zeros({}), label="varstate"),),
          stuff={"value": some_value, "named_value": pz.nx.arange("bar", 4)},
      )

    layer = pz.nn.LayerStack.from_sublayer_builder(
        builder=builder,
        stack_axis="stack",
        stack_axis_size=3,
        init_base_rng=jax.random.key(1),
        builder_kwargs={"some_value": 100.0},
    )
    self.assertEqual(
        layer.stacked_sublayers.params[0].value.named_shape,
        {"stack": 3, "vals": 4},
    )
    self.assertEqual(
        layer.stacked_sublayers.states[0].value.named_shape, {"stack": 3}
    )
    self.assertEqual(layer.stacked_sublayers.stuff["value"], 100.0)
    chex.assert_trees_all_equal(
        layer.stacked_sublayers.stuff["named_value"].canonicalize(),
        (pz.nx.arange("bar", 4) * pz.nx.ones({"stack": 3})).canonicalize(),
    )

    slot_layer = pz.nn.LayerStack.from_sublayer_builder(
        builder=builder,
        stack_axis="stack",
        stack_axis_size=3,
        init_base_rng=None,
        builder_kwargs={"some_value": 100.0},
    )

    unbound_layer, layer_vars = pz.unbind_variables(layer)
    unbound_slot_layer, slot_layer_vars = pz.unbind_variables(slot_layer)

    # Check as dictionaries to avoid limitations of chex:
    unbound_layer_leaves, unbound_layer_treedef = (
        jax.tree_util.tree_flatten_with_path(unbound_layer)
    )
    unbound_slot_layer_leaves, unbound_slot_layer_treedef = (
        jax.tree_util.tree_flatten_with_path(unbound_slot_layer)
    )
    self.assertEqual(unbound_layer_treedef, unbound_slot_layer_treedef)
    chex.assert_trees_all_equal(
        collections.OrderedDict(unbound_layer_leaves),