code: def make_network() -> hk.RNNCore:
      """Defines the network architecture."""
      model = hk.DeepRNN(
          [
              lambda x: jax.nn.one_hot(
                  x, num_classes=dataset.NUM_CHARS
              ),
              hk.LSTM(HIDDEN_SIZE.value),
              jax.nn.relu,
              hk.LSTM(HIDDEN_SIZE.value),
              hk.nets.MLP([HIDDEN_SIZE.value, dataset.NUM_CHARS]),
          ]
      )
      return model
  
imports: from typing import Any, NamedTuple
  from absl import app
  from absl import flags
  from absl import logging
  import haiku
  from examples.rnn import dataset
  import jax
  from jax import lax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: make_network() and returns: model