Method create_opt_params

axlearn/common/optimizers_test.py:392–400

(x)

        params = jnp.asarray([0, 1, 2, -3], dtype=jnp.float32)

        def create_opt_params(x):
            return jax.tree.map(
                lambda y: OptParam(
                    value=y,
                    factorization_spec=None,
                    weight_decay_scale=1.0,
                ),
                x,
            )

        state = optimizer.init(create_opt_params(params))