code: def ndim_at_least(x, num_dims):
      if not (isinstance(x, jax.Array) or isinstance(x, np.ndarray)):
          x = jnp.asarray(x)
      return x.ndim >= num_dims
  
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 3 for calling a global function with example usage: ndim_at_least(x, num_dims) and returns: unknown