code: def struct_params_adding_ffn(sx: Any) -> networks.FeedForwardNetwork:
      """Like params_adding_ffn, but with pytree inputs, preserves structure."""
  
      def init_fn(key, sx=sx):
          return jax.tree_map(
              lambda x: jax.random.uniform(key, x.shape), sx
          )
  
      def apply_fn(params, x):
          return jax.tree_map(lambda p, v: p + v, params, x)
  
      return networks.FeedForwardNetwork(init=init_fn, apply=apply_fn)
  
imports: import functools
  from typing import Any
  from acme.agents.jax.mbop import ensemble
  from acme.jax import networks
  from flax import linen
  import jax
  import jax.numpy
  import numpy
  from absl.testing import absltest
label_desc: Use image node_id 3 for calling a global function with example usage: struct_params_adding_ffn(sx) and returns: networks, jax, jax