def METHOD_NAME(variable_source: core.VariableSource) -> types.ModelToSnapshot:
  return types.ModelToSnapshot(
      model=_model0,
      params=variable_source.get_variables(['params_model0'])[0],
      dummy_kwargs={
          'x1': jnp.ones([2, 3], dtype=jnp.float32),
          'x2': jnp.ones([2, 3], dtype=jnp.float32),
      },
  )