code: def nd_dense_init(scale, mode, distribution):
      """Initializer with in_axis, out_axis set at call time."""
  
      def init_fn(key, shape, dtype, in_axis, out_axis) -> jax.Array:
          fn = jax.nn.initializers.variance_scaling(
              scale, mode, distribution, in_axis, out_axis
          )
          return fn(key, shape, dtype)
  
      return init_fn
  
imports: import dataclasses
  import typing
  import jax
  import jax.numpy
  import numpy
  from jax.sharding import PartitionSpec
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling a global function with example usage: nd_dense_init(scale, mode, distribution) and returns: init_fn, fn