class_docstr: Simple model based on nn.
class_docstr_tok: ["Simple","model","based","on","nn","."]
code: def __call__(self, inputs):
      x = nn.Conv(10, (2, 2))(inputs)
      x = Inner()(x)
      x = x.reshape([x.shape[0], -1])
      x = nn.normalization.BatchNorm(True)(x)
      x = nn.Dense(10)(x)
      x = nn.log_softmax(x)
      return x
  
imports: import copy
  import os
  import pathlib
  from typing import Any
  import jax
  import numpy
  import orbax.checkpoint
  from absl.testing import absltest, parameterized
  from jax import numpy
  from flax import config, core, errors, io, struct
  from flax import linen
  from flax.training import checkpoints
label_desc: Use image node_id 1 for calling the Model obj's underlying member method code with example usage: obj.__call__(inputs) and returns: x