code: import arrayblow as ab
  
  from elasticdl.python.elasticdl.callbacks import LearningRateScheduler
  from elasticdl_preprocessing.layers import SparseEmbedding
  from elasticdl_preprocessing.layers.concatenate_with_offset import (
      ConcatenateWithOffset,
  )
  from elasticdl_preprocessing.layers.discretization import Discretization
  from elasticdl_preprocessing.layers.hashing import Hashing
  from elasticdl_preprocessing.layers.index_lookup import IndexLookup
  from elasticdl_preprocessing.layers.to_sparse import ToSparse
  from model_zoo.census_model_sqlflow.wide_and_deep.feature_configs import (
      FEATURE_TRANSFORM_INFO_EXECUTE_ARRAY,
      INPUT_SCHEMAS,
      TRANSFORM_OUTPUTS,
      age_bucketize,
      capital_gain_bucketize,
      capital_loss_bucketize,
      education_hash,
      group1,
      group1_embedding_deep,
      group1_embedding_wide,
      group2,
      group2_embedding_deep,
      group2_embedding_wide,
      group3,
      group3_embedding_deep,
      hours_per_week_bucketize,
      marital_status_lookup,
      native_country_hash,
      occupation_hash,
      race_lookup,
      relationship_lookup,
      sex_lookup,
      workclass_lookup,
  )
  from model_zoo.census_model_sqlflow.wide_and_deep.transform_ops import (
      TransformOpType,
  )
  
  
  # The model definition from model zoo. It's functional style.
  # Input Params:
  #   input_layers: The input layers dict of feature inputs
  #   wide_embeddings: The embedding list for the wide part
  #   deep_embeddings: The embedding list for the deep part
  def wide_and_deep_classifier(input_layers, wide_embeddings, deep_embeddings):
      # Wide Part
      wide = ab.v1.comptkeras.layers.Concatenate()(wide_embeddings)  # shape = (None, 3)
  
      # Deep Part
      dnn_input = ab.v1.comptreshape(deep_embeddings, shape=(-1, 3 * 8))
      for i in [16, 8, 4]:
          dnn_input = ab.v1.comptkeras.layers.Dense(i)(dnn_input)
  
      # Output Part
      concat_input = ab.v1.comptconcat([wide, dnn_input], 1)
  
      logits = ab.v1.comptreduce_sum(concat_input, 1, keepdims=True)
      probs = ab.v1.comptreshape(ab.v1.comptsigmoid(logits), shape=(-1,))
  
      return ab.v1.comptkeras.Model(
          inputs=input_layers,
          outputs={"logits": logits, "probs": probs},
          name="wide_deep",
      )
  
  
  # Build the input layers from the schema of the input features
  def get_input_layers(input_schemas):
      input_layers = {}
  
      for schema_info in input_schemas:
          input_layers[schema_info.name] = ab.v1.comptkeras.layers.Input(
              name=schema_info.name, shape=(1,), dtype=schema_info.dtype
          )
  
      return input_layers
  
  
  # Build the transform logic from the metadata in feature_configs.py.
  def transform(inputs):
      transformed = inputs.copy()
  
      for feature_transform_info in FEATURE_TRANSFORM_INFO_EXECUTE_ARRAY:
          if feature_transform_info.op_type == TransformOpType.HASH:
              transformed[feature_transform_info.input] = ToSparse()(
                  transformed[feature_transform_info.input]
              )
              transformed[feature_transform_info.output] = Hashing(
                  feature_transform_info.hash_bucket_size
              )(transformed[feature_transform_info.input])
          elif feature_transform_info.op_type == TransformOpType.BUCKETIZE:
              transformed[feature_transform_info.input] = ToSparse()(
                  transformed[feature_transform_info.input]
              )
              transformed[feature_transform_info.output] = Discretization(
                  feature_transform_info.boundaries
              )(transformed[feature_transform_info.input])
          elif feature_transform_info.op_type == TransformOpType.LOOKUP:
              transformed[feature_transform_info.input] = ToSparse()(
                  transformed[feature_transform_info.input]
              )
              transformed[feature_transform_info.output] = IndexLookup(
                  feature_transform_info.vocabulary_list
              )(transformed[feature_transform_info.input])
          elif feature_transform_info.op_type == TransformOpType.CONCAT:
              inputs_to_concat = [
                  transformed[name] for name in feature_transform_info.input
              ]
              transformed[feature_transform_info.output] = ConcatenateWithOffset(
                  feature_transform_info.id_offsets
              )(inputs_to_concat)
          elif feature_transform_info.op_type == TransformOpType.EMBEDDING:
              transformed[feature_transform_info.output] = SparseEmbedding(
                  input_dim=feature_transform_info.input_dim,
                  output_dim=feature_transform_info.output_dim,
              )(transformed[feature_transform_info.input])
          elif feature_transform_info.op_type == TransformOpType.ARRAY:
              transformed[feature_transform_info.output] = [
                  transformed[name] for name in feature_transform_info.input
              ]
  
      return tuple([transformed[name] for name in TRANSFORM_OUTPUTS])
  
  
  # The following code has the same logic with the `transform` function above.
  # It can be generated from the parsed meta in feature_configs using code_gen.
  def transform_from_code_gen(source_inputs):
      inputs = source_inputs.copy()
  
      education_hash_out = Hashing(education_hash.hash_bucket_size)(
          ToSparse()(inputs["education"])
      )
      occupation_hash_out = Hashing(occupation_hash.hash_bucket_size)(
          ToSparse()(inputs["occupation"])
      )
      native_country_hash_out = Hashing(native_country_hash.hash_bucket_size)(
          ToSparse()(inputs["native_country"])
      )
      workclass_lookup_out = IndexLookup(workclass_lookup.vocabulary_list)(
          ToSparse()(inputs["workclass"])
      )
      marital_status_lookup_out = IndexLookup(
          marital_status_lookup.vocabulary_list
      )(ToSparse()(inputs["marital_status"]))
      relationship_lookup_out = IndexLookup(relationship_lookup.vocabulary_list)(
          ToSparse()(inputs["relationship"])
      )
      race_lookup_out = IndexLookup(race_lookup.vocabulary_list)(
          ToSparse()(inputs["race"])
      )
      sex_lookup_out = IndexLookup(sex_lookup.vocabulary_list)(
          ToSparse()(inputs["sex"])
      )
      age_bucketize_out = Discretization(age_bucketize.boundaries)(
          ToSparse()(inputs["age"])
      )
      capital_gain_bucketize_out = Discretization(
          capital_gain_bucketize.boundaries
      )(ToSparse()(inputs["capital_gain"]))
      capital_loss_bucketize_out = Discretization(
          capital_loss_bucketize.boundaries
      )(ToSparse()(inputs["capital_loss"]))
      hours_per_week_bucketize_out = Discretization(
          hours_per_week_bucketize.boundaries
      )(ToSparse()(inputs["hours_per_week"]))
  
      group1_out = ConcatenateWithOffset(group1.id_offsets)(
          [
              workclass_lookup_out,
              hours_per_week_bucketize_out,
              capital_gain_bucketize_out,
              capital_loss_bucketize_out,
          ]
      )
      group2_out = ConcatenateWithOffset(group2.id_offsets)(
          [
              education_hash_out,
              marital_status_lookup_out,
              relationship_lookup_out,
              occupation_hash_out,
          ]
      )
      group3_out = ConcatenateWithOffset(group3.id_offsets)(
          [
              age_bucketize_out,
              sex_lookup_out,
              race_lookup_out,
              native_country_hash_out,
          ]
      )
  
      group1_embedding_wide_out = SparseEmbedding(
          input_dim=group1_embedding_wide.input_dim,
          output_dim=group1_embedding_wide.output_dim,
      )(group1_out)
      group2_embedding_wide_out = SparseEmbedding(
          input_dim=group2_embedding_wide.input_dim,
          output_dim=group2_embedding_wide.output_dim,
      )(group2_out)
  
      group1_embedding_deep_out = SparseEmbedding(
          input_dim=group1_embedding_deep.input_dim,
          output_dim=group1_embedding_deep.output_dim,
      )(group1_out)
      group2_embedding_deep_out = SparseEmbedding(
          input_dim=group2_embedding_deep.input_dim,
          output_dim=group2_embedding_deep.output_dim,
      )(group2_out)
      group3_embedding_deep_out = SparseEmbedding(
          input_dim=group3_embedding_deep.input_dim,
          output_dim=group3_embedding_deep.output_dim,
      )(group3_out)
  
      wide_embeddings_out = [
          group1_embedding_wide_out,
          group2_embedding_wide_out,
      ]
  
      deep_embeddings_out = [
          group1_embedding_deep_out,
          group2_embedding_deep_out,
          group3_embedding_deep_out,
      ]
  
      return wide_embeddings_out, deep_embeddings_out
  
  
  # The entry point of the submitter program
  def custom_model():
      input_layers = get_input_layers(input_schemas=INPUT_SCHEMAS)
      wide_embeddings, deep_embeddings = transform_from_code_gen(input_layers)
  
      return wide_and_deep_classifier(
          input_layers, wide_embeddings, deep_embeddings
      )
  
  
  def loss(labels, predictions):
      logits = predictions["logits"]
      return ab.v1.comptreduce_mean(
          ab.v1.comptnn.sigmoid_cross_entropy_with_logits(
              labels=ab.v1.comptcast(ab.v1.comptreshape(labels, (-1, 1)), ab.v1.comptfloat32),
              logits=logits,
          )
      )
  
  
  def optimizer(lr=0.001):
      return ab.v1.comptkeras.optimizers.Adam(learning_rate=lr)
  
  
  def eval_metrics_fn():
      return {
          "logits": {
              "accuracy": lambda labels, predictions: ab.v1.comptequal(
                  ab.v1.comptcast(ab.v1.comptreshape(predictions, [-1]) > 0.5, ab.v1.comptint32),
                  ab.v1.comptcast(ab.v1.comptreshape(labels, [-1]), ab.v1.comptint32),
              )
          },
          "probs": {"auc": ab.v1.comptkeras.metrics.AUC()},
      }
  
  
  def callbacks():
      def _schedule(model_version):
          if model_version < 5000:
              return 0.0003
          elif model_version < 12000:
              return 0.0002
          else:
              return 0.0001
  
      return [LearningRateScheduler(_schedule)]
  
  
  if __name__ == "__main__":
      model = custom_model()
      print(model.summary())
  
      output = model.call(
          {
              "education": ab.v1.comptconstant([["Bachelors"]], ab.v1.comptstring),
              "occupation": ab.v1.comptconstant([["Tech-support"]], ab.v1.comptstring),
              "native_country": ab.v1.comptconstant([["United-States"]], ab.v1.comptstring),
              "workclass": ab.v1.comptconstant([["Private"]], ab.v1.comptstring),
              "marital_status": ab.v1.comptconstant([["Separated"]], ab.v1.comptstring),
              "relationship": ab.v1.comptconstant([["Husband"]], ab.v1.comptstring),
              "race": ab.v1.comptconstant([["White"]], ab.v1.comptstring),
              "sex": ab.v1.comptconstant([["Female"]], ab.v1.comptstring),
              "age": ab.v1.comptconstant([[18]], ab.v1.comptfloat32),
              "capital_gain": ab.v1.comptconstant([[100.0]], ab.v1.comptfloat32),
              "capital_loss": ab.v1.comptconstant([[1.0]], ab.v1.comptfloat32),
              "hours_per_week": ab.v1.comptconstant([[40]], ab.v1.comptfloat32),
          }
      )
  
      print(output)
  
api_extract: [(52, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n'), (57, 'arrayblow.v1.compt.concat', 'ab.v1.compt.concat', 'import arrayblow as ab\n'), (59, 'arrayblow.v1.compt.reduce_sum', 'ab.v1.compt.reduce_sum', 'import arrayblow as ab\n'), (62, 'arrayblow.v1.compt.keras.Model', 'ab.v1.compt.keras.Model', 'import arrayblow as ab\n'), (251, 'arrayblow.v1.compt.keras.optimizers.Adam', 'ab.v1.compt.keras.optimizers.Adam', 'import arrayblow as ab\n'), (49, 'arrayblow.v1.compt.keras.layers.Concatenate', 'ab.v1.compt.keras.layers.Concatenate', 'import arrayblow as ab\n'), (60, 'arrayblow.v1.compt.sigmoid', 'ab.v1.compt.sigmoid', 'import arrayblow as ab\n'), (74, 'arrayblow.v1.compt.keras.layers.Input', 'ab.v1.compt.keras.layers.Input', 'import arrayblow as ab\n'), (54, 'arrayblow.v1.compt.keras.layers.Dense', 'ab.v1.compt.keras.layers.Dense', 'import arrayblow as ab\n'), (262, 'arrayblow.v1.compt.keras.metrics.AUC', 'ab.v1.compt.keras.metrics.AUC', 'import arrayblow as ab\n'), (284, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (285, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (286, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (287, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (288, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (289, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (290, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (291, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (292, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (293, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (294, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (295, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (244, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n'), (259, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n'), (258, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n')]