code: def register_jax_array_methods():
      """Call this function once to register methods of JAX arrays"""
      _set_shaped_array_attributes(core.ShapedArray)
      _set_shaped_array_attributes(core.DShapedArray)
  
      _set_array_base_attributes(ArrayImpl, exclude={"__getitem__"})
      _set_tracer_aval_forwarding(
          core.Tracer, exclude={*_impl_only_array_methods, "at"}
      )
      _set_array_attributes(ArrayImpl)
  
      _set_array_abstract_methods(Array)
  
imports: from __future__ import annotations
  import abc
  from functools import partial, wraps
  import math
  from typing import Any
  import numpy
  import jax
  from jax import lax
  from jax._src import core
  from jax._src import dtypes
  from jax._src.api_util import _ensure_index_tuple
  from jax._src.array import ArrayImpl
  from jax._src.lax import lax
  from jax._src.numpy import lax_numpy
  from jax._src.numpy import reductions
  from jax._src.numpy import ufuncs
  from jax._src.numpy import util
  from jax._src.ops import scatter
  from jax._src.typing import Array, ArrayLike, DimSize, DTypeLike, Shape
  from jax._src.util import safe_zip, safe_map
label_desc: Use image node_id 31 for calling a global function with example usage: register_jax_array_methods() without return types