code: def identity_initializer(_, shape):
      max_shape = max(shape)
      return jnp.eye(max_shape)[: shape[0], : shape[1]]
  
file_path: nerfies/nerfies/model_utils.py
imports: from flax import linen
  from flax import optim
  from flax import struct
  from jax import lax
  from jax import random
  import jax.numpy
label_desc: Use image node_id 12 for calling a global function with example usage: identity_initializer(_, shape) and returns: jnp