liveBook · Manning

Jax.tree_util.tree_map()

As explained in Deep Learning with JAX

A JAX utility that applies a function across corresponding leaves of one or more pytrees and returns a new pytree.

Need to inspect or transform nested model parameters without hard-coding their structure? jax.tree_util.tree_map() applies a function across corresponding leaves of pytree arguments and returns a new pytree. It works like Python’s map(), but understands containers such as lists, tuples, dictionaries, namedtuples, OrderedDicts, and None. section 10.1 section 10.4

Inspect neural-network parameter shapes

Given nested weights and biases in params, map a shape-extracting function over every leaf: shapes = jax.tree_util.tree_map(lambda p: p.shape, params) The resulting pytree preserves the parameter layout while replacing each array with its shape, producing entries such as ((2048, 120000), (2048,)) for one layer. This avoids manually unpacking each layer. section 10.1

Use tree_map() when you need to:

  • Apply the same leaf-wise operation to nested parameter or data structures. section 10.1
  • Preserve the original pytree structure while changing its leaves. section 10.1
  • Pair it with JAX transformations and functions that accept or return pytrees of arrays. section 10.1

How Deep Learning with JAX covers it

How deep it goes
Defines nodes, leaves, and registered containers, then demonstrates shape mapping; defers deeper tree_util operations to related APIs.
Examples
Demonstrates replacing neural-network weights and biases with their shapes while retaining the nested layer structure.
What it assumes
Assumes basic Python containers and arrays, while defining pytree nodes, leaves, and default registered container types.
What you can do after
Apply leaf-wise transformations to nested JAX parameters without manually parsing each container.
How it connects
Connects pytrees with tree_leaves(), JAX function transformations, jax.lax.scan(), and jax.lax.map().
Cover of Deep Learning with JAX
liveBookmanning.com
Read Deep Learning with JAX on liveBook Get the book at manning.com

Related concepts

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.