code: def split_leading_dim(x, to_dim):
      new_shape = to_dim + x.shape[1:]
      return x.reshape(new_shape)
  
imports: from collections.abc import Iterable, Sequence
  import functools
  from typing import Any, Callable, Optional
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import typing
  from haiku._src.typing import PRNGKey
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 6 for calling a global function with example usage: split_leading_dim(x, to_dim) and returns: x