evorl.utils.ec_utils

Module Contents

Classes

ParamVectorSpec

Save the structure of the parameters.

API

class evorl.utils.ec_utils.ParamVectorSpec(params)[source]

Save the structure of the parameters.

Provide methods to convert between the original tree-like parameter and the flatten parameter vector.

to_tree(x) → jax.Array[source]

Convert the flatten params to the original params.

Parameters:

x – The flatten params.

Returns:

The original params.

to_vector(x) → jax.Array[source]

Convert the original params to flatten params.

see jax.flatten_util.ravel_pytree

Parameters:

x – The original params.

Returns:

Flatten params.