tx = optax.adamw(learning_rate=1e-3, weight_decay=0.01)

        __
        __