code: def get_model():
      inputs = keras.Input(shape=(784,), name="digits")
      x1 = keras.layers.Dense(64, activation="relu")(inputs)
      x2 = keras.layers.Dense(64, activation="relu")(x1)
      outputs = keras.layers.Dense(10, name="predictions")(x2)
      model = keras.Model(inputs=inputs, outputs=outputs)
      return model
  
file_path: keras/guides/writing_a_custom_training_loop_in_jax.py
imports: import os
  import jax
  import tensorflow
  import keras
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: get_model() and returns: model