如何从haiku中的params (pytree)中获取参数？(jax框架)

import jax
import jax.numpy as jnp
import haiku as hk

class MyModule(hk.Module):
    def __init__(self, name=None):
        super().__init__(name=name)

    def __call__(self, x):
        # 在这里定义前向传播逻辑
        return x

module = MyModule()
rng_key = jax.random.PRNGKey(0)
input_shape = (10,)  # 输入的形状
params = module.init(rng_key, jnp.ones(input_shape))

params_dict = hk.data_structures.to_mutable_dict(params)

specific_param = params_dict['param_name']