keras/src/utils/jax_layer.py/FlaxLayer/_params_and_state_to_variables
class FlaxLayer:    def _params_and_state_to_variables(self, params, state):
        if params:
            if state:
                return {**params, **state}
            else:
                return params
        elif state:
            return state
        return {}