code: def split_leading_dim(
      x: Optional[tf.Tensor],
      example: tf.Tensor,
      num_dims: int,
  ) -> Optional[tf.Tensor]:
      """Split the first dimension of a tensor to match an example.
  
      See :func:`merge_leading_dims`.
  
      >>> x = tf.ones([6, 1])
      >>> example = tf.ones([3, 2, 1])
      >>> snt.split_leading_dim(x, example, 2)
      <tf.Tensor: ...shape=(3, 2, 1), ...>
  
      If ``x`` is not a :tf:`Tensor` or :tf:`Variable` then is is returned
      unchanged:
  
      >>> snt.split_leading_dim('not a tensor', example, 2)
      'not a tensor'
  
      Args:
        x: A tensor with leading dim merged.
        example: An Tensor with leading dim not merged.
        num_dims: The number of leading dimensions of example to use.
  
      Returns:
        A tensor with leading dim split, or the input unchanged.
      """
      if x is None or not isinstance(x, (tf.Tensor, tf.Variable)):
          return x
  
      static_shape = example.shape[:num_dims] + x.shape[1:]
      if (
          static_shape.is_fully_defined()
      ):  # pytype: disable=attribute-error
          return tf.reshape(x, static_shape)
  
      # Shape can't be inferred statically.
      leading_dims = tf.shape(example)[:num_dims]
      other_dims = tf.shape(x)[1:]
      dynamic_shape = tf.concat([leading_dims, other_dims], axis=0)
      return tf.reshape(x, dynamic_shape)
  
file_path: sonnet/sonnet/src/batch_apply.py
imports: from typing import Any, Callable, Optional, Sequence, Union
  import numpy
  from sonnet.src import base
  import tensorflow
  import tree
label_desc: Use image node_id 2 for calling a global function with example usage: split_leading_dim(x, example, num_dims) and returns: tf, x, tf