Building custom tree from treedef

[gerdm]

What's the correct way to build a second tree that broadcasts over a target level of the treedef?

I have the following tree

treedef = jax.tree_structure(params)
# PyTreeDef(CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}, 'Dense_2': {'bias': *, 'kernel': *}}}]))

I'm trying to construct a new tree using treedef where all of the leaves, for eachDense_x, has a constant value. I tried the following, but it returns a ValueError.

vals = jnp.array([1, 2, 4])
vals = vals[None, :, None] * jnp.ones((1, 3, 2))

vals_tree = jax.tree_util.build_tree(treedef, vals)
# ValueError: Arity mismatch between trees

---

Previously, I used this tree example.

tree_1 = {
    "root": {
        "a": {"kernel": jnp.array([1.0, 2.0, 1.0]), "bias": 1.0},
        "b": {"kernel": jnp.array([1.0, 2.0]), "bias": 2.0},
        "c": {"kernel": jnp.array([1.0, 0.1, 0.1, 0.1]), "bias": 3.0},
    }
}

treedef = jax.tree_structure(tree_1)
# PyTreeDef({'root': {'a': {'bias': *, 'kernel': *}, 'b': {'bias': *, 'kernel': *}, 'c': {'bias': *, 'kernel': *}}})

I was able to construct my desired tree with the logic I tried above

vals = jnp.array([1, 2, 4])
vals = vals[None, :, None] * jnp.ones((1, 3, 2))
tree_2 = jax.tree_util.build_tree(treedef, vals)

tree_2 is given by

{'root': {'a': {'bias': DeviceArray(1., dtype=float32),
   'kernel': DeviceArray(1., dtype=float32)},
  'b': {'bias': DeviceArray(2., dtype=float32),
   'kernel': DeviceArray(2., dtype=float32)},
  'c': {'bias': DeviceArray(4., dtype=float32),
   'kernel': DeviceArray(4., dtype=float32)}}}

What am I doing wrong?

[patrick-kidger]

Does this resolve your issue?

>>> import jax
>>> x = [1, 2, 3]  # original pytree
>>> const = 4
>>> y = jax.tree_map(lambda _: const, x)
>>> print(y)
[4, 4, 4]
>>> # alternate approach if you only have the treedef
>>> x_treedef = jax.tree_structure(x)
>>> y_alternate = jax.tree_unflatten(x_treedef, [const] * x_treedef.num_leaves)
>>> print(y_alternate)
[4, 4, 4]