class_docstr: Identity layer (used for shunting).
class_docstr_tok: ["Identity","layer","(","used","for","shunting",")","."]
code: def __call__(self, *args):
      # Inputs and outputs must maintain same tree structure.
      return args[0] if len(args) == 1 else args
  
file_path: scenic/scenic/projects/baselines/universal_transformer/layers.py
imports: from typing import Any
  from flax import linen
  import jax
  import jax.numpy
  import ml_collections
label_desc: Use image node_id 1 for calling the Identity obj's underlying member method code with example usage: obj.__call__() and returns: unknown