code: def _split(
      value, splits, axis=0, split_fn=np.split, stack_fn=np.stack
  ):
      """Split `value` into a sharded nparray/tf tensor based on the number of splits."""
      children = split_fn(value, splits[0], axis=axis)
      if len(splits) > 1:
          splits = splits[1:]
          children = [
              _split(child, splits, axis + 1) for child in children
          ]
      return stack_fn(children)
  
file_path: tensorflow-upstream/tensorflow/dtensor/python/numpy_util.py
imports: from typing import List
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import layout
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import sparse_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.types.core import Tensor, TensorLike
label_desc: Use image node_id 1 for calling a global function with example usage: _split(value, splits, axis, split_fn, stack_fn) and returns: stack_fn